Function robot_body_pos_b

TextOpTracker/scripts/deploy_mujoco.py:472–502

Robot body positions in body frame

(sim_data, motion_loader, t)

def robot_body_pos_b(sim_data, motion_loader, t):
    """Robot body positions in body frame"""
    # Get robot anchor pose
    robot_anchor_pos = sim_data.body(motion_loader.anchor_body_name).xpos.copy().reshape(1, 3)
    robot_anchor_quat = sim_data.body(motion_loader.anchor_body_name).xquat.copy().reshape(1, 4)

    # Get all robot body poses
    robot_body_positions = []
    robot_body_orientations = []

    for body_name in motion_loader.body_names:
        body_pos = sim_data.body(body_name).xpos.copy().reshape(1, 3)
        body_quat = sim_data.body(body_name).xquat.copy().reshape(1, 4)
        robot_body_positions.append(body_pos)
        robot_body_orientations.append(body_quat)

    robot_body_pos = np.concatenate(robot_body_positions, axis=0)  # [14, 3]
    robot_body_quat = np.concatenate(robot_body_orientations, axis=0)  # [14, 4]

    # Transform to body frame - expand robot_anchor to match robot_body shape
    robot_anchor_pos_expanded = robot_anchor_pos.repeat(len(motion_loader.body_names), axis=0)  # [14, 3]
    robot_anchor_quat_expanded = robot_anchor_quat.repeat(len(motion_loader.body_names), axis=0)  # [14, 4]

    pos_b, _ = subtract_frame_transforms(
        torch.from_numpy(robot_anchor_pos_expanded).float(),
        torch.from_numpy(robot_anchor_quat_expanded).float(),
        torch.from_numpy(robot_body_pos).float(),
        torch.from_numpy(robot_body_quat).float(),
    )

    return np.array(pos_b, dtype=np.float32).flatten()  # [42]


def robot_body_ori_b(sim_data, motion_loader, t):