unilab.envs.motion_tracking.g1.motion_loader

Motion loading and sampling for motion tracking tasks.

Classes

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.envs.motion_tracking.g1.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:
class unilab.envs.motion_tracking.g1.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

get_clip_end_frames(frame_idx)[source]

Return the inclusive global end frame for each indexed clip.

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.envs.motion_tracking.g1.motion_loader.MotionSampler[source]

Bases: object

Handles motion frame sampling with different strategies.

Parameters:

  • motion_loader (MotionLoader)

  • mode (Literal['start', 'clip_start', 'uniform', 'adaptive', 'mixed'])

  • num_envs (int)

  • bin_count (int | None)

  • adaptive_lambda (float)

  • adaptive_kernel_size (int)

  • adaptive_uniform_ratio (float)

  • adaptive_alpha (float)

  • start_ratio (float)

__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)[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.

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()[source]

Advance all frames by one step.

get_current_motion(out=None)[source]

Get motion data at current frames for all environments.

Parameters:

out (MotionData | None)

Return type:

MotionData