fix(loss-func): dheading loss ignores train/val mode prefix (+1 more)

[andrewwhitecdw]

Small fixes in traj_pred/models/loss_func.py:

fix: dheading loss ignores train/val mode prefix

Fix: Replace:

def compute_dheading(data, specs):
    local_traj = data[f'train_out_local_traj_tp']
    local_heading_vec = local_traj[..., -2:]
    heading = vec_to_heading(local_heading_vec)
    mse = heading.pow(2).mean()
    return mse

with:

def compute_dheading(data, specs):
    mode = specs.get('mode', 'train')
    local_traj = data[f'{mode}_out_local_traj_tp']
    local_heading_vec = local_traj[..., -2:]
    heading = vec_to_heading(local_heading_vec)
    mse = heading.pow(2).mean()
    return mse

fix: local_orient_heading loss ignores train/val mode prefix

Fix: Replace:

def compute_local_orient_heading(data, specs):
    local_traj = data[f'train_out_local_traj_tp']
    local_orient = local_traj[..., 3:-2]
    if local_orient.shape[-1] == 6:
        local_orient = rot6d_to_quat(local_orient)
    heading = get_heading(local_orient)
    mse = heading.pow(2).mean()
    return mse

with:

def compute_local_orient_heading(data, specs):
    mode = specs.get('mode', 'train')
    local_traj = data[f'{mode}_out_local_traj_tp']
    local_orient = local_traj[..., 3:-2]
    if local_orient.shape[-1] == 6:
        local_orient = rot6d_to_quat(local_orient)
    heading = get_heading(local_orient)
    mse = heading.pow(2).mean()
    return mse

Files changed

• traj_pred/models/loss_func.py

[andrewwhitecdw]

Closing this sweep-generated PR: sole commit is missing a valid Signed-off-by trailer. It does not meet the sweep requirements (single signed-off commit).