unilab.tasks.motion_tracking.common.motion_loader

Shared motion loading and sampling for motion-tracking tasks.

Functions

compute_frame_blend(times, duration, ...)

Compute frame indices and blend weights for interpolation.

compute_motion_velocities(base_poss, ...)

Numerically differentiate a trajectory into base/dof velocities.

interpolate_motion(base_poss_input, ...)

Resample a root+joint trajectory from input_fps to output_fps.

quat_slerp(q1, q2, t)

Spherical linear interpolation between two quaternions (wxyz format).

Classes

InterpolatedMotion

Root/joint trajectory resampled to the output frame rate.

MotionData

Container for motion data at specific frame(s).

MotionLoader

Loads and provides access to motion data from NPZ files.

MotionSampler

Handles motion frame sampling with different strategies.
class unilab.tasks.motion_tracking.common.motion_loader.MotionData[source]

Bases: object

Container for motion data at specific frame(s).

Parameters:
joint_pos: ndarray
joint_vel: ndarray
body_pos_w: ndarray
body_quat_w: ndarray
body_lin_vel_w: ndarray
body_ang_vel_w: ndarray
__init__(joint_pos, joint_vel, body_pos_w, body_quat_w, body_lin_vel_w, body_ang_vel_w)
Parameters:
unilab.tasks.motion_tracking.common.motion_loader.quat_slerp(q1, q2, t)[source]

Spherical linear interpolation between two quaternions (wxyz format).

The computation runs in the input dtype; pass float64 arrays for a float64 interpolation path.

Parameters:
Return type:

ndarray

class unilab.tasks.motion_tracking.common.motion_loader.InterpolatedMotion[source]

Bases: object

Root/joint trajectory resampled to the output frame rate.

Parameters:
output_frames: int
base_poss: ndarray
base_rots: ndarray
dof_poss: ndarray
base_lin_vels: ndarray
base_ang_vels: ndarray
dof_vels: ndarray
__init__(output_frames, base_poss, base_rots, dof_poss, base_lin_vels, base_ang_vels, dof_vels)
Parameters:
unilab.tasks.motion_tracking.common.motion_loader.compute_motion_velocities(base_poss, base_rots, dof_poss, dt)[source]

Numerically differentiate a trajectory into base/dof velocities.

Parameters:
Return type:

tuple[ndarray, ndarray, ndarray]

unilab.tasks.motion_tracking.common.motion_loader.compute_frame_blend(times, duration, input_frames)[source]

Compute frame indices and blend weights for interpolation.

Parameters:
Return type:

tuple[ndarray, ndarray, ndarray]

unilab.tasks.motion_tracking.common.motion_loader.interpolate_motion(base_poss_input, base_rots_input, dof_poss_input, *, input_fps, output_fps)[source]

Resample a root+joint trajectory from input_fps to output_fps.

Positions and joint angles are linearly interpolated, root quaternions (wxyz) use slerp, and velocities are computed by numerical differentiation at the output rate.

Parameters:
Return type:

InterpolatedMotion

class unilab.tasks.motion_tracking.common.motion_loader.MotionLoader[source]

Bases: object

Loads and provides access to motion data from NPZ files.

Parameters:
__init__(motion_file, body_indices=None)[source]

Initialize motion loader.

Parameters:

  • motion_file (str | Sequence[str]) – Path to one NPZ file, or a sequence of NPZ files

  • body_indices (ndarray | None) – Optional indices into the NPZ body axis. The exported motion files currently keep MuJoCo body-id layout, so these indices are expected to follow that convention.

get_clip_indices(frame_idx)[source]

Map global frame indices to clip indices.

Parameters:

frame_idx (ndarray)

Return type:

ndarray

make_motion_data_buffer(num_frames)[source]

Allocate a reusable MotionData buffer for frame-index gathers.

Parameters:

num_frames (int)

Return type:

MotionData

get_motion_at_frame(frame_idx, out=None)[source]

Get motion data at specified frame indices.

Parameters:

  • frame_idx (ndarray) – Frame indices (N,)

  • out (MotionData | None) – Optional reusable output buffer.

Return type:

MotionData

Returns:

MotionData at specified frames

class unilab.tasks.motion_tracking.common.motion_loader.MotionSampler[source]

Bases: object

Handles motion frame sampling with different strategies.

Parameters:
__init__(motion_loader, mode, num_envs, bin_count=None, adaptive_lambda=0.8, adaptive_kernel_size=1, adaptive_uniform_ratio=0.1, adaptive_alpha=0.001, start_ratio=0.0, rng=None)[source]

Initialize motion sampler.

Parameters:

  • motion_loader (MotionLoader) – Motion loader instance

  • mode (Literal['start', 'clip_start', 'uniform', 'adaptive', 'mixed']) – Sampling mode (“start”, “clip_start”, “uniform”, “adaptive”, “mixed”)

  • num_envs (int) – Number of parallel environments

  • bin_count (int | None) – Number of bins for adaptive sampling (auto if None)

  • adaptive_lambda (float) – Decay factor for adaptive kernel

  • adaptive_kernel_size (int) – Kernel size for adaptive sampling

  • adaptive_uniform_ratio (float) – Uniform sampling ratio for adaptive mode

  • adaptive_alpha (float) – EMA alpha for failure count updates

  • start_ratio (float) – Fraction of envs forced to frame 0 in “mixed” mode (remaining envs are uniformly sampled). Lets buffer concentrate launch-transition samples while keeping motion-clip coverage.

  • rng (Generator | None)

sample_frames(env_ids)[source]

Sample motion frames for specified environments.

Parameters:

env_ids (ndarray) – Environment indices to sample for

Return type:

ndarray

Returns:

Sampled frame indices

update_failure_stats(terminated, current_frames=None)[source]

Update failure statistics for adaptive sampling.

Parameters:

  • terminated (ndarray) – Boolean array indicating which environments terminated

  • current_frames (ndarray | None) – Optional current frame indices (uses internal if None)

step(env_ids=None)[source]

Advance selected frames by one step and return their clip-end rows.

Parameters:

env_ids (ndarray | None)

Return type:

ndarray

get_current_motion(out=None)[source]

Get motion data at current frames for all environments.

Parameters:

out (MotionData | None)

Return type:

MotionData