unilab.tasks.motion_tracking.common.kernels

Parallel CPU kernels for the fixed motion-tracking hot-term set.

Functions

configure_motion_kernel_runtime()

Initialize the task-local Numba worker mask once on the cold path.

reward_motion_body_ang_vel_kernel(reference, ...)

Write the global body-angular-velocity exponential reward.

reward_motion_body_lin_vel_kernel(reference, ...)

Write the global body-linear-velocity exponential reward.

reward_motion_body_ori_kernel(reference, ...)

Write the relative body-orientation exponential reward.

reward_motion_body_pos_kernel(reference, ...)

Write the relative body-position exponential reward.

termination_anchor_pos_kernel(...)

Write the per-environment anchor-height termination mask.

update_motion_metrics_kernel(env_ids, ...)

Write all MotionCommand error metrics for the selected environment rows.

update_motion_relative_state_kernel(env_ids, ...)

Write all MotionCommand relative transforms for selected rows.

update_object_relative_state_kernel(env_ids, ...)

Write BoxMotionCommand object pose and velocity for selected rows.

unilab.tasks.motion_tracking.common.kernels.configure_motion_kernel_runtime()[source]

Initialize the task-local Numba worker mask once on the cold path.

Return type:

None

unilab.tasks.motion_tracking.common.kernels.reward_motion_body_ang_vel_kernel(reference, actual, body_ids, std, out)[source]

Write the global body-angular-velocity exponential reward.

Parameters:
Return type:

None

unilab.tasks.motion_tracking.common.kernels.reward_motion_body_lin_vel_kernel(reference, actual, body_ids, std, out)[source]

Write the global body-linear-velocity exponential reward.

Parameters:
Return type:

None

unilab.tasks.motion_tracking.common.kernels.reward_motion_body_ori_kernel(reference, actual, body_ids, std, out)[source]

Write the relative body-orientation exponential reward.

Parameters:
Return type:

None

unilab.tasks.motion_tracking.common.kernels.reward_motion_body_pos_kernel(reference, actual, body_ids, std, out)[source]

Write the relative body-position exponential reward.

Parameters:
Return type:

None

unilab.tasks.motion_tracking.common.kernels.termination_anchor_pos_kernel(motion_body_pos_w, robot_body_pos_w, anchor_body_idx, threshold, out)[source]

Write the per-environment anchor-height termination mask.

Parameters:
Return type:

None

unilab.tasks.motion_tracking.common.kernels.update_motion_metrics_kernel(env_ids, anchor_body_idx, motion_body_pos_w, robot_body_pos_w, motion_body_quat_w, robot_body_quat_w, motion_body_lin_vel_w, robot_body_lin_vel_w, motion_body_ang_vel_w, robot_body_ang_vel_w, body_pos_relative_w, body_quat_relative_w, motion_joint_pos, robot_joint_pos, motion_joint_vel, robot_joint_vel, error_anchor_pos, error_anchor_rot, error_anchor_lin_vel, error_anchor_ang_vel, error_body_pos, error_body_rot, error_body_lin_vel, error_body_ang_vel, error_joint_pos, error_joint_vel)[source]

Write all MotionCommand error metrics for the selected environment rows.

env_ids is a cold-path-owned all-row index buffer for normal steps and the reset row list for partial resets. Keeping one kernel for both paths avoids a second production mathematical implementation while preserving the manager’s row-scoped reset contract.

Parameters:
Return type:

None

unilab.tasks.motion_tracking.common.kernels.update_motion_relative_state_kernel(env_ids, anchor_body_idx, motion_body_pos_local_w, motion_body_pos_w, motion_body_quat_w, robot_body_pos_w, robot_body_quat_w, body_pos_relative_w, body_quat_relative_w, motion_anchor_pos_b, motion_anchor_ori_b, robot_body_pos_b, robot_body_ori_b)[source]

Write all MotionCommand relative transforms for selected rows.

Parameters:
Return type:

None

unilab.tasks.motion_tracking.common.kernels.update_object_relative_state_kernel(env_ids, robot_anchor_pos_w, robot_anchor_quat_w, object_pos_w, object_quat_w, object_lin_vel_w, object_state_b)[source]

Write BoxMotionCommand object pose and velocity for selected rows.

Parameters:
Return type:

None