Skip to content

Loss Functions

losses

Imitation learning losses: invariant, geodesic, hybrid.

Functions

hybrid_invariant_pose_loss 

hybrid_invariant_pose_loss(
    pred_positions,
    pred_quaternions,
    demo_positions,
    demo_quaternions,
    pred_invariants=None,
    demo_invariants=None,
    alpha=0.5,
    beta=1.0,
)

alpha * invariant_loss + (1-alpha) * pose_loss. If pred_invariants/demo_invariants are None, only pose loss is used (alpha ignored for invariant part).

Source code in src/dhb_xr/losses/hybrid_loss.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def hybrid_invariant_pose_loss(
    pred_positions: np.ndarray,
    pred_quaternions: np.ndarray,
    demo_positions: np.ndarray,
    demo_quaternions: np.ndarray,
    pred_invariants: Optional[np.ndarray] = None,
    demo_invariants: Optional[np.ndarray] = None,
    alpha: float = 0.5,
    beta: float = 1.0,
) -> float:
    """
    alpha * invariant_loss + (1-alpha) * pose_loss.
    If pred_invariants/demo_invariants are None, only pose loss is used (alpha ignored for invariant part).
    """
    loss_pose = 0.0
    n = len(pred_positions)
    assert n == len(demo_positions) and n == len(pred_quaternions) and n == len(demo_quaternions)
    for i in range(n):
        loss_pose += se3_geodesic_loss_np(
            pred_positions[i], pred_quaternions[i],
            demo_positions[i], demo_quaternions[i],
            beta=beta,
        )
    if pred_invariants is not None and demo_invariants is not None:
        loss_inv = invariant_matching_loss(pred_invariants, demo_invariants)
        return float(alpha * loss_inv + (1 - alpha) * loss_pose)
    return float(loss_pose)

invariant_matching_loss 

invariant_matching_loss(pred_inv, demo_inv, method='dhb_dr', weights=None)

pred_inv, demo_inv: (N, 2*k). method 'dhb_dr' (Euler) or 'dhb_qr' (quaternion). For dhb_qr, angular quaternion part uses geodesic; else L2 with optional angle wrap.

Source code in src/dhb_xr/losses/invariant_loss.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def invariant_matching_loss(
    pred_inv: np.ndarray,
    demo_inv: np.ndarray,
    method: str = "dhb_dr",
    weights: Optional[np.ndarray] = None,
) -> float:
    """
    pred_inv, demo_inv: (N, 2*k). method 'dhb_dr' (Euler) or 'dhb_qr' (quaternion).
    For dhb_qr, angular quaternion part uses geodesic; else L2 with optional angle wrap.
    """
    pred_inv = np.asarray(pred_inv)
    demo_inv = np.asarray(demo_inv)
    assert pred_inv.shape == demo_inv.shape
    k = pred_inv.shape[1] // 2
    if weights is None:
        weights = np.ones(pred_inv.shape[1])
    if method == "dhb_qr":
        m_lin = np.sum(weights[0] * (pred_inv[:, 0] - demo_inv[:, 0]) ** 2)
        q_lin = pred_inv[:, 1:5]
        q_lin_d = demo_inv[:, 1:5]
        m_ang = np.sum(weights[k] * (pred_inv[:, k] - demo_inv[:, k]) ** 2)
        q_ang = pred_inv[:, k + 1 : k + 5]
        q_ang_d = demo_inv[:, k + 1 : k + 5]
        loss_lin_q = quaternion_geodesic_loss_np(q_lin, q_lin_d)
        loss_ang_q = quaternion_geodesic_loss_np(q_ang, q_ang_d)
        return float(m_lin + m_ang + loss_lin_q + loss_ang_q)
    diff = pred_inv - demo_inv
    return float(np.sum(weights * (diff ** 2)))

so3_geodesic_loss 

so3_geodesic_loss(R_pred, R_demo)

Dispatch to numpy or torch.

Source code in src/dhb_xr/losses/geodesic_loss.py 
47
48
49
50
51
def so3_geodesic_loss(R_pred, R_demo):
    """Dispatch to numpy or torch."""
    if hasattr(R_pred, "numpy"):
        return so3_geodesic_loss_torch(R_pred, R_demo)
    return so3_geodesic_loss_np(R_pred, R_demo)

Overview

Loss functions for imitation learning with DHB invariants.

Main Functions

invariant_matching_loss

invariant_loss

Imitation learning losses in invariant space.

Functions

invariant_matching_loss 

invariant_matching_loss(pred_inv, demo_inv, method='dhb_dr', weights=None)

pred_inv, demo_inv: (N, 2*k). method 'dhb_dr' (Euler) or 'dhb_qr' (quaternion). For dhb_qr, angular quaternion part uses geodesic; else L2 with optional angle wrap.

Source code in src/dhb_xr/losses/invariant_loss.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def invariant_matching_loss(
    pred_inv: np.ndarray,
    demo_inv: np.ndarray,
    method: str = "dhb_dr",
    weights: Optional[np.ndarray] = None,
) -> float:
    """
    pred_inv, demo_inv: (N, 2*k). method 'dhb_dr' (Euler) or 'dhb_qr' (quaternion).
    For dhb_qr, angular quaternion part uses geodesic; else L2 with optional angle wrap.
    """
    pred_inv = np.asarray(pred_inv)
    demo_inv = np.asarray(demo_inv)
    assert pred_inv.shape == demo_inv.shape
    k = pred_inv.shape[1] // 2
    if weights is None:
        weights = np.ones(pred_inv.shape[1])
    if method == "dhb_qr":
        m_lin = np.sum(weights[0] * (pred_inv[:, 0] - demo_inv[:, 0]) ** 2)
        q_lin = pred_inv[:, 1:5]
        q_lin_d = demo_inv[:, 1:5]
        m_ang = np.sum(weights[k] * (pred_inv[:, k] - demo_inv[:, k]) ** 2)
        q_ang = pred_inv[:, k + 1 : k + 5]
        q_ang_d = demo_inv[:, k + 1 : k + 5]
        loss_lin_q = quaternion_geodesic_loss_np(q_lin, q_lin_d)
        loss_ang_q = quaternion_geodesic_loss_np(q_ang, q_ang_d)
        return float(m_lin + m_ang + loss_lin_q + loss_ang_q)
    diff = pred_inv - demo_inv
    return float(np.sum(weights * (diff ** 2)))

quaternion_geodesic_loss_np 

quaternion_geodesic_loss_np(q1, q2)

Sum of squared quaternion geodesic distances. q1, q2: (N, 4) wxyz.

Source code in src/dhb_xr/losses/invariant_loss.py 
16
17
18
19
20
def quaternion_geodesic_loss_np(q1: np.ndarray, q2: np.ndarray) -> float:
    """Sum of squared quaternion geodesic distances. q1, q2: (N, 4) wxyz."""
    dot = np.abs(np.sum(q1 * q2, axis=-1))
    dot = np.clip(dot, 0, 1)
    return np.sum((2 * np.arccos(dot)) ** 2)

se3_geodesic_loss

geodesic_loss

SO(3) and SE(3) geodesic losses.

Functions

se3_geodesic_loss_np 

se3_geodesic_loss_np(pos_pred, quat_pred, pos_demo, quat_demo, beta=1.0)

Position L2 + beta * SO3 geodesic. Quaternions wxyz.

Source code in src/dhb_xr/losses/geodesic_loss.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def se3_geodesic_loss_np(
    pos_pred: np.ndarray,
    quat_pred: np.ndarray,
    pos_demo: np.ndarray,
    quat_demo: np.ndarray,
    beta: float = 1.0,
) -> float:
    """Position L2 + beta * SO3 geodesic. Quaternions wxyz."""
    loss_pos = np.sum((pos_pred - pos_demo) ** 2)
    R_pred = geom.quat_to_rot(quat_pred)
    R_demo = geom.quat_to_rot(quat_demo)
    if R_pred.ndim == 2:
        R_pred = R_pred.reshape(1, 3, 3)
        R_demo = R_demo.reshape(1, 3, 3)
    R_diff = np.einsum("...ji,...jk->...ik", R_pred, R_demo)
    rvec = np.array([geom.rot_to_axis_angle(R_diff[i]) for i in range(len(R_diff))])
    loss_rot = np.sum(rvec ** 2)
    return float(loss_pos + beta * loss_rot)

so3_geodesic_loss 

so3_geodesic_loss(R_pred, R_demo)

Dispatch to numpy or torch.

Source code in src/dhb_xr/losses/geodesic_loss.py 
47
48
49
50
51
def so3_geodesic_loss(R_pred, R_demo):
    """Dispatch to numpy or torch."""
    if hasattr(R_pred, "numpy"):
        return so3_geodesic_loss_torch(R_pred, R_demo)
    return so3_geodesic_loss_np(R_pred, R_demo)

so3_geodesic_loss_np 

so3_geodesic_loss_np(R_pred, R_demo)

||Log(R_pred^T R_demo)||^2. R_pred, R_demo: (3, 3) or (N, 3, 3).

Source code in src/dhb_xr/losses/geodesic_loss.py 
15
16
17
18
19
20
21
22
23
24
def so3_geodesic_loss_np(R_pred: np.ndarray, R_demo: np.ndarray) -> float:
    """||Log(R_pred^T R_demo)||^2. R_pred, R_demo: (3, 3) or (N, 3, 3)."""
    R_pred = np.asarray(R_pred)
    R_demo = np.asarray(R_demo)
    if R_pred.ndim == 2:
        R_pred = R_pred.reshape(1, 3, 3)
        R_demo = R_demo.reshape(1, 3, 3)
    R_diff = np.einsum("...ji,...jk->...ik", R_pred, R_demo)
    rvec = np.array([geom.rot_to_axis_angle(R_diff[i]) for i in range(len(R_diff))])
    return float(np.sum(rvec ** 2))

hybrid_invariant_pose_loss

hybrid_loss

Hybrid invariant + pose-space loss for imitation learning.

Functions

hybrid_invariant_pose_loss 

hybrid_invariant_pose_loss(
    pred_positions,
    pred_quaternions,
    demo_positions,
    demo_quaternions,
    pred_invariants=None,
    demo_invariants=None,
    alpha=0.5,
    beta=1.0,
)

alpha * invariant_loss + (1-alpha) * pose_loss. If pred_invariants/demo_invariants are None, only pose loss is used (alpha ignored for invariant part).

Source code in src/dhb_xr/losses/hybrid_loss.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def hybrid_invariant_pose_loss(
    pred_positions: np.ndarray,
    pred_quaternions: np.ndarray,
    demo_positions: np.ndarray,
    demo_quaternions: np.ndarray,
    pred_invariants: Optional[np.ndarray] = None,
    demo_invariants: Optional[np.ndarray] = None,
    alpha: float = 0.5,
    beta: float = 1.0,
) -> float:
    """
    alpha * invariant_loss + (1-alpha) * pose_loss.
    If pred_invariants/demo_invariants are None, only pose loss is used (alpha ignored for invariant part).
    """
    loss_pose = 0.0
    n = len(pred_positions)
    assert n == len(demo_positions) and n == len(pred_quaternions) and n == len(demo_quaternions)
    for i in range(n):
        loss_pose += se3_geodesic_loss_np(
            pred_positions[i], pred_quaternions[i],
            demo_positions[i], demo_quaternions[i],
            beta=beta,
        )
    if pred_invariants is not None and demo_invariants is not None:
        loss_inv = invariant_matching_loss(pred_invariants, demo_invariants)
        return float(alpha * loss_inv + (1 - alpha) * loss_pose)
    return float(loss_pose)

Usage Example

from dhb_xr.losses.invariant_loss import invariant_matching_loss
from dhb_xr.encoder.dhb_dr import encode_dhb_dr

# Encode demonstrations and predictions
demo_inv = encode_dhb_dr(demo_pos, demo_quat)
pred_inv = encode_dhb_dr(pred_pos, pred_quat)

# Compute invariant matching loss
loss = invariant_matching_loss(
    demo_inv['linear_motion_invariants'],
    demo_inv['angular_motion_invariants'],
    pred_inv['linear_motion_invariants'],
    pred_inv['angular_motion_invariants']
)