Motion Reference

Motion tracking tasks score a robot against reference motion. Eden serves that motion along two paths, and this page documents the classes of both. The motion tracking guide explains the data schema, the multi-actor pattern and which tracking term reads which field.

• The manager path — MimicRLEnvBase owns a MotionManager whose MotionTerm loads a dataset, reads the frame at floor((start + t) * fps) (no interpolation at lookup), serves a window of future frames and runs the motion-id and start-time curricula. Tracking terms read it as env.motion_manager.data[<entity_name>]. This is the training path.
• The command path — MotionCommand (a command term) plays .npz clips verbatim on a plain RLEnvBase, for frozen-policy inference and single-clip tracking. It is documented with the command terms.

Both expose the same world-frame vocabulary: links_*_w for the tracked links and base_*_w for the floating base. Quaternions are wxyz.

Motion manager and MotionTerm

Motion manager and MotionTerm base: motion-reference data for imitation tasks.

A :class:MotionTerm behaves like a PyTorch Dataset of motion clips (some generate references on the fly), and the :class:MotionManager fills per-env reference buffers each step for :class:~eden.envs.base.MimicRLEnvBase.

Under DDP, each rank is assigned only a subset of trajectories (data_source and the _local_to_global map shrink per rank) while the motion-weight and outcome buffers stay sized to the full motion-id space, so cross-rank all_reduce stays shape-symmetric.

Attributes

MOTION_TERM_REGISTRY module-attribute

MOTION_TERM_REGISTRY = TermRegistry('MOTION_TERM', MotionTerm)

Classes

MotionTerm

MotionTerm(env: EnvBase, options: MotionTermOptions)

Bases: ConfigurableTermBase[MotionTermOptions]

Base class for motion terms, which serve a similar role to Dataset in PyTorch.

Some of the motion terms may return the motion reference by generating it on the fly.

PARAMETER	DESCRIPTION
dataset_options	The options for the dataset.
entities_map	The mapping from the entity name in the environment to the entity name in the dataset. - key: entity name in the environment - value: motion sequence key in the dataset
target_dofs_name	The names of the degrees of freedom to target.
target_links_name	The names of the links to target.
env_motion_groups	Per-env clip groups keyed on a heterogeneous entity's variant; None samples every env from the whole library. See :class:~eden.options.managers.motions.MotionGroupOptions. TYPE: MotionGroupOptions | None

Source code in eden/managers/motion_manager.py

def __init__(
    self,
    env: EnvBase,
    options: MotionTermOptions,
) -> None:
    super().__init__(env=env, options=options)

    assert options.dataset_options is not None, "dataset_options is required for motion terms."
    # NOTE: setup dataset
    self.data_source = DATASET_REGISTRY.get(options.dataset_options.name)(
        options.dataset_options,
        device=self.device,
    )

    for motion_sequence_key in self.entities_map.values():
        if motion_sequence_key not in self.data_source.motion_sequence_keys:
            raise ValueError(
                f"Motion sequence key {motion_sequence_key} not found in the dataset. Available keys: {self.data_source.motion_sequence_keys}"
            )

    # NOTE: resolve the indices of the motion sequence keys
    target_dofs_idx: dict[str, list[int]] = {}
    target_links_idx: dict[str, list[int]] = {}
    for key in self.data_source.motion_sequence_keys:
        indices, names = resolve_matching_names(
            self.target_dofs_name[key],
            self.data_source.motion_dofs_name[key],
            preserve_order=True,
        )
        self.target_dofs_name[key] = names
        # Reorder on any non-identity selection, not only a subset: a full-length permuted target list
        # must permute the columns too, or every name lookup reads the wrong DOF.
        if indices != list(range(len(self.data_source.motion_dofs_name[key]))):
            target_dofs_idx[key] = indices
        indices, names = resolve_matching_names(
            self.target_links_name[key],
            self.data_source.motion_links_name[key],
            preserve_order=True,
        )
        self.target_links_name[key] = names
        if indices != list(range(len(self.data_source.motion_links_name[key]))):
            target_links_idx[key] = indices

    self.data_source.select_dofs(target_dofs_idx)
    self.data_source.select_links(target_links_idx)

    motion_weights_source: str | None = None
    if isinstance(self.motion_weights, str):
        motion_weights_source = self.motion_weights
        self.motion_weights = load_text(motion_weights_source).get("motion_weights", None)
        if self.motion_weights is None:
            raise ValueError(f"`motion_weights` key not found in {motion_weights_source}")
    if self.motion_weights is not None:
        if len(self.motion_weights) != len(self.data_source):
            # _motion_weights and the per-motion outcome buffers index a
            # shared GLOBAL motion-id space; a length mismatch with the
            # dataset corrupts every later read. Catch it at construction
            # rather than on the first sampling call.
            src = f" (loaded from {motion_weights_source!r})" if motion_weights_source else ""
            raise ValueError(
                f"`motion_weights`{src} has length {len(self.motion_weights)}, but the "
                f"dataset has {len(self.data_source)} trajectories. Each entry in "
                f"`motion_weights` is the initial sampling weight for one motion id; "
                f"supply exactly one weight per loaded trajectory."
            )
        self._motion_weights = torch.tensor(
            self.motion_weights,
            device=self.device,
            dtype=torch.float,
        )
    else:
        self._motion_weights = torch.ones(
            len(self.data_source),
            device=self.device,
            dtype=torch.float,
        )

    # Pin the global motion-id space at construction time, before any per-rank
    # subsetting via enable_trajectories(). _motion_weights and the curriculum
    # outcome buffers are kept at this size on every rank — that is what lets
    # the iter-end dist.all_reduce shape-match across ranks even when each
    # rank only loads its assigned trajectory slice into VRAM.
    self._global_num_motions: int = len(self.data_source)
    self._local_to_global: torch.Tensor = torch.arange(
        self._global_num_motions, device=self.device, dtype=torch.long
    )

    # Per-env clip groups (``env_motion_groups``): each clip's group id in the GLOBAL motion-id space, resolved
    # here from the dataset's per-clip labels; the per-rank mask and the per-env group follow in
    # ``enable_trajectories`` / ``set_envs_idx_assignments``. All three stay ``None`` when the option is unset.
    self._global_motion_group: torch.Tensor | None = None
    self._group_motion_mask: torch.Tensor | None = None
    self._env_group: torch.Tensor | None = None
    if self.env_motion_groups is not None:
        self._global_motion_group = self._resolve_global_motion_groups(self.env_motion_groups)

    # Temporal start-time bin weights + adaptive curriculum live in an
    # AdaptiveBinSampler collaborator, (re)created in enable_trajectories when
    # motion_bin_length_s is set (class-level _bin_sampler=None default until then).
    if self.adaptive_sampling and self.motion_bin_length_s is None:
        raise ValueError("adaptive_sampling=True requires motion_bin_length_s to be set.")
    if self.compute_per_motion_metrics and not self.adaptive_motion_weights:
        # Per-motion metric drain piggybacks on the curriculum's DDP
        # cadence and ``all_reduce_fn`` to share one synced collective
        # per PPO iter. Without the curriculum nothing would drain the
        # metric buffers, so fail loudly rather than let them grow.
        raise ValueError(
            "compute_per_motion_metrics=True requires adaptive_motion_weights=True; "
            "the metric drain shares the curriculum's DDP cadence."
        )

    # Per-motion-id outcome curriculum (hard-motion mining), orthogonal to
    # the bin-level adaptive_sampling. Lives in MotionOutcomeCurriculum,
    # which owns the success/failure counts + summary caches and rewrites
    # ``_motion_weights`` from drained outcome rates. ``None`` when the
    # curriculum is off, so disabled terms pay nothing.
    self._curriculum: MotionOutcomeCurriculum | None = (
        MotionOutcomeCurriculum(self) if self.adaptive_motion_weights else None
    )

    # Per-motion tracking-error metrics live in a PerMotionMetrics
    # collaborator, created in ``set_envs_idx_assignments`` once
    # ``num_assigned_envs`` is known (and only when ``compute_per_motion_metrics``
    # is set, so disabled terms pay nothing). ``None`` until then / when disabled.
    self._metrics: PerMotionMetrics | None = None

Attributes

dataset_options class-attribute instance-attribute

dataset_options: DatasetOptions | None = None

entities_map class-attribute instance-attribute

entities_map: dict[str, str] = {}

target_dofs_name class-attribute instance-attribute

target_dofs_name: dict[str, str | list[str]] = {}

target_links_name class-attribute instance-attribute

target_links_name: dict[str, str | list[str]] = {}

motion_weights class-attribute instance-attribute

motion_weights: str | list[float] | None = None

env_motion_groups class-attribute instance-attribute

env_motion_groups: MotionGroupOptions | None = None

env_starting_stub_sampling_strategy class-attribute instance-attribute

env_starting_stub_sampling_strategy: Literal['independent', 'concat_motion_bins'] = 'independent'

motion_start_from_middle_range class-attribute instance-attribute

motion_start_from_middle_range: tuple[float, float] | list[float] = (0.0, 0.0)

motion_bin_length_s class-attribute instance-attribute

motion_bin_length_s: float | None = None

adaptive_sampling class-attribute instance-attribute

adaptive_sampling: bool = False

adaptive_kernel_size class-attribute instance-attribute

adaptive_kernel_size: int = 1

adaptive_lambda class-attribute instance-attribute

adaptive_lambda: float = 0.8

adaptive_uniform_ratio class-attribute instance-attribute

adaptive_uniform_ratio: float = 0.1

adaptive_alpha class-attribute instance-attribute

adaptive_alpha: float = 0.001

adaptive_failure_rate_max_over_mean class-attribute instance-attribute

adaptive_failure_rate_max_over_mean: float = 10.0

adaptive_motion_weights class-attribute instance-attribute

adaptive_motion_weights: bool = False

motion_weight_success_discount class-attribute instance-attribute

motion_weight_success_discount: float = 0.999

motion_weight_failure_amplify class-attribute instance-attribute

motion_weight_failure_amplify: float = 2.0

motion_weight_floor class-attribute instance-attribute

motion_weight_floor: float = 0.001

motion_weight_update_mode class-attribute instance-attribute

motion_weight_update_mode: Literal['compound', 'snapshot'] = 'snapshot'

motion_weight_uniform_mix class-attribute instance-attribute

motion_weight_uniform_mix: float = 0.0

motion_weight_failure_rate_max_over_mean class-attribute instance-attribute

motion_weight_failure_rate_max_over_mean: float = 200.0

motion_weight_update_every_n_iters class-attribute instance-attribute

motion_weight_update_every_n_iters: int = 1

start_at_timestep_zero_prob class-attribute instance-attribute

start_at_timestep_zero_prob: float = 0.0

freeze_at_timestep_zero_prob class-attribute instance-attribute

freeze_at_timestep_zero_prob: float = 0.0

compute_per_motion_metrics class-attribute instance-attribute

compute_per_motion_metrics: bool = False

metric_anchor_link_name class-attribute instance-attribute

metric_anchor_link_name: str | None = None

log_per_motion_metrics class-attribute instance-attribute

log_per_motion_metrics: bool = False

log_motion_id_to_filename class-attribute instance-attribute

log_motion_id_to_filename: bool = False

data_source instance-attribute

data_source: DatasetBase = DATASET_REGISTRY.get(options.dataset_options.name)(options.dataset_options, device=self.device)

num_trajectories cached property

num_trajectories: int

Return the number of trajectories in the motion buffer.

For a generative motion buffer, this just returns a placeholder number.

num_assigned_envs property

num_assigned_envs: int

The number of envs assigned to this motion buffer.

complete_motion_lengths property

complete_motion_lengths: Tensor

The full motion reference lengths for each assigned env. the unit is in seconds.

assigned_motion_lengths property

assigned_motion_lengths: Tensor

The motion reference lengths for each assigned env. the unit is in seconds.

Methods:

set_envs_idx_assignments

set_envs_idx_assignments(envs_idx: slice) -> None

Assign the environment indices this motion buffer is responsible for.

The motion buffer is expected to fill the motion reference data for the assigned envs only.

PARAMETER	DESCRIPTION
envs_idx	The environment IDs to assign to this motion buffer. TYPE: slice

Source code in eden/managers/motion_manager.py

def set_envs_idx_assignments(self, envs_idx: slice) -> None:
    """Assign the environment indices this motion buffer is responsible for.

    The motion buffer is expected to fill the motion reference data for the assigned envs only.

    Parameters
    ----------
    envs_idx : slice
        The environment IDs to assign to this motion buffer.
    """
    assert envs_idx.step is None, "Currently only support continuous env IDs."
    self.assigned_env_slice = envs_idx

    self._assigned_env_motion_selection = torch.zeros(
        self.num_assigned_envs,
        dtype=torch.long,
        device=self.device,
    )
    self._motion_buffer_start_time_s = torch.zeros(
        self.num_assigned_envs,
        dtype=torch.float,
        device=self.device,
    )

    assigned_ids = self.assigned_ids_from_envs_idx(None)
    if self.env_starting_stub_sampling_strategy == "concat_motion_bins" and self.motion_bin_length_s is None:
        raise ValueError(
            "env_starting_stub_sampling_strategy='concat_motion_bins' requires motion_bin_length_s "
            "to be set so that motion bin weights are initialized. Got motion_bin_length_s=None."
        )
    if self._global_motion_group is not None:
        self._env_group = self._resolve_env_groups(self.env_motion_groups)
    if self.env_starting_stub_sampling_strategy == "independent":
        self._sample_assigned_env_motion_selection(assigned_ids)
        self._sample_env_motion_start_time(assigned_ids)
    elif self.env_starting_stub_sampling_strategy == "concat_motion_bins":
        self._sample_as_concat_motion_bins(assigned_ids)

    # Allocate the per-motion metric collaborator now that num_assigned_envs
    # is known. Gated on the term option so disabled terms pay nothing.
    if self.compute_per_motion_metrics:
        self._metrics = PerMotionMetrics(self)

enable_trajectories

enable_trajectories(traj_ids: Tensor | slice | None = None) -> None

Enable only the selected motion reference trajectories.

Assuming all motion reference trajectories are stored in a sequential way, the motion manager will call this function to enable only the selected trajectories, depending on how multi-processing is implemented. By "enabled", it means the motion buffer will generate the motion data for only the enabled trajectories.

_motion_weights and the per-motion outcome buffers stay sized to the full pre-subset motion-id space on every rank — only data_source and the new _local_to_global map shrink. This is what lets DDP all_reduce stay shape-symmetric while each rank holds only its assigned trajectory slice in VRAM.

PARAMETER	DESCRIPTION
traj_ids	The trajectory IDs to enable. start from 0 to self.num_trajectories - 1 in this motion buffer. If None, enable all trajectories. Out-of-range values are ignored, so no motion is enabled if every id is out of range. TYPE: Tensor | slice | None DEFAULT: None

Source code in eden/managers/motion_manager.py

def enable_trajectories(self, traj_ids: torch.Tensor | slice | None = None) -> None:
    """Enable only the selected motion reference trajectories.

    Assuming all motion reference trajectories are stored in a sequential way, the motion manager will call this
    function to enable only the selected trajectories, depending on how multi-processing is implemented.
    By "enabled", it means the motion buffer will generate the motion data for only the enabled trajectories.

    ``_motion_weights`` and the per-motion outcome buffers stay sized to the
    full pre-subset motion-id space on every rank — only ``data_source`` and
    the new ``_local_to_global`` map shrink. This is what lets DDP
    ``all_reduce`` stay shape-symmetric while each rank holds only its
    assigned trajectory slice in VRAM.

    Parameters
    ----------
    traj_ids : torch.Tensor | slice | None, optional
        The trajectory IDs to enable. start from 0 to `self.num_trajectories - 1` in this motion buffer.
        If None, enable all trajectories.
        Out-of-range values are ignored, so no motion is enabled if every id is out of range.
    """
    if traj_ids is not None:
        # ``traj_ids`` are LOCAL ids relative to the term's *current*
        # ``num_trajectories`` (per the docstring). Compose them with the
        # existing ``_local_to_global`` BEFORE we shrink ``data_source`` so
        # repeated calls remain correct (e.g. global=[3, 7, 9] then a
        # local=[0, 2] subset must produce global=[3, 9], not [0, 2]).
        if isinstance(traj_ids, slice):
            start, stop, step = traj_ids.indices(len(self._local_to_global))
            local_idx = torch.arange(start, stop, step, device=self.device, dtype=torch.long)
        else:
            local_idx = torch.as_tensor(traj_ids, device=self.device, dtype=torch.long)
        self._local_to_global = self._local_to_global[local_idx]
        self.data_source.select_trajectories(traj_ids)
        # invalidate cached num_trajectories since the dataset has been subsetted
        self.__dict__.pop("num_trajectories", None)
    # (Re)create the bin sampler from the now-current (possibly subsetted)
    # data_source; its weights are per-motion so they must track the library.
    self._bin_sampler = AdaptiveBinSampler(self) if self.motion_bin_length_s is not None else None
    if self._global_motion_group is not None:
        # (G, M_local): which of this rank's clips each group may draw. Built once per library change, so the
        # per-reset draw is a single masked multinomial with no host sync.
        local_group = self._global_motion_group[self._local_to_global]
        num_groups = len(self._group_names)
        self._group_motion_mask = local_group.unsqueeze(0) == torch.arange(
            num_groups, device=self.device
        ).unsqueeze(1)

rewind_to_start

rewind_to_start(envs_idx: Tensor) -> None

Rewind playback to frame 0 of each env's currently-assigned clip (no resampling).

Zeros this term's per-env start offset. The per-term half of :meth:MotionManager.rewind_to_start; public so callers needn't poke _motion_buffer_start_time_s (allocated in :meth:set_envs_idx_assignments at manager construction, so it always exists by the time playback is rewound).

Source code in eden/managers/motion_manager.py

def rewind_to_start(self, envs_idx: torch.Tensor) -> None:
    """Rewind playback to frame 0 of each env's currently-assigned clip (no resampling).

    Zeros this term's per-env start offset. The per-term half of
    :meth:`MotionManager.rewind_to_start`; public so callers needn't poke
    ``_motion_buffer_start_time_s`` (allocated in :meth:`set_envs_idx_assignments` at
    manager construction, so it always exists by the time playback is rewound).
    """
    self._motion_buffer_start_time_s[envs_idx] = 0.0

reset

reset(envs_idx: Tensor | None, symmetric_augmentation_mask_buffer: Tensor | None = None) -> None

Reset the motion buffer for the specified envs_idx.

PARAMETER	DESCRIPTION
envs_idx	The environment IDs to reset the motion buffer, which is among the entire num_envs. If None, reset all buffers. TYPE: Tensor | None
symmetric_augmentation_mask_buffer	The buffer to store the symmetric augmentation mask for each env. The motion buffer is expected to sample and fill the symmetric augmentation mask for each env. Because some of the motions do or do not support symmetric augmentation. NOTE: the motion buffer does not apply the symmetric augmentation; it only updates the mask. TYPE: Tensor | None DEFAULT: None

Source code in eden/managers/motion_manager.py

def reset(
    self,
    envs_idx: torch.Tensor | None,
    symmetric_augmentation_mask_buffer: torch.Tensor | None = None,
) -> None:
    """
    Reset the motion buffer for the specified envs_idx.

    Parameters
    ----------
    envs_idx : torch.Tensor | None
        The environment IDs to reset the motion buffer, which is among the entire num_envs.
        If None, reset all buffers.
    symmetric_augmentation_mask_buffer : torch.Tensor | None
        The buffer to store the symmetric augmentation mask for each env. The
        motion buffer is expected to sample and fill the symmetric augmentation mask for each env. Because some
        of the motions do or do not support symmetric augmentation.
        NOTE: the motion buffer does not apply the symmetric augmentation; it only updates the mask.
    """
    self._sample_assigned_env_starting_stub(envs_idx)

    if symmetric_augmentation_mask_buffer is not None:
        if envs_idx is not None and not isinstance(envs_idx, torch.Tensor):
            envs_idx = torch.as_tensor(envs_idx, device=self.device)
        if is_env_mask(envs_idx):
            # Over-sample full buffer, mask-write to avoid GPU→CPU sync
            rand_mask = torch.rand(len(symmetric_augmentation_mask_buffer), device=self.device) < 0.5
            symmetric_augmentation_mask_buffer[envs_idx] = rand_mask[envs_idx]
        else:
            num_envs = self.num_assigned_envs if envs_idx is None else len(envs_idx)
            symmetric_augmentation_mask_buffer[envs_idx] = torch.rand(num_envs, device=self.device) < 0.5

fill_init_reference_state

fill_init_reference_state(motion_sequence_key: str, envs_idx: Sequence[int] | Tensor, env_origins: Tensor | None, state_buffer: MotionReferenceState) -> None

Fill the initial reference state into the buffer for the specified envs only.

PARAMETER	DESCRIPTION
motion_sequence_key	The motion sequence key to fill the initial reference state from. TYPE: str
envs_idx	The environment IDs to fill the initial reference state. TYPE: Sequence[int] | Tensor
env_origins	The origins for each robot/env, which is also affecting the motion reference data, shape (num_envs, 3). None: every env's origin is the world origin. TYPE: Tensor | None
state_buffer	The MotionReferenceState to fill the initial reference state, which is a buffer for all envs, shape (num_envs, ...). TYPE: MotionReferenceState

Source code in eden/managers/motion_manager.py

def fill_init_reference_state(
    self,
    motion_sequence_key: str,
    envs_idx: Sequence[int] | torch.Tensor,
    env_origins: torch.Tensor | None,
    state_buffer: MotionReferenceState,
) -> None:
    """Fill the initial reference state into the buffer for the specified envs only.

    Parameters
    ----------
    motion_sequence_key : str
        The motion sequence key to fill the initial reference state from.
    envs_idx : Sequence[int] | torch.Tensor
        The environment IDs to fill the initial reference state.
    env_origins : torch.Tensor | None
        The origins for each robot/env, which is also affecting the motion reference data, shape (num_envs, 3).
        None: every env's origin is the world origin.
    state_buffer : MotionReferenceState
        The MotionReferenceState to fill the initial reference state, which is a buffer for all envs,
        shape (num_envs, ...).
    """
    assigned_ids = self.envs_idx_to_assigned_ids(envs_idx)
    data = self.data_source.motion_sequences[motion_sequence_key]
    motion_ids = self._assigned_env_motion_selection[assigned_ids]

    # Sampled start times are clamped to (buffer_length - 1) / framerate (see
    # ``sample_start_times`` / ``_sample_as_concat_motion_bins``), so the frame
    # index is in-range by construction; the explicit min() guards the float
    # round-trip of that clamp (start * framerate landing exactly on
    # buffer_length), holding the last frame instead of reading past the clip —
    # which the masked CBT path would zero-fill: the zero-quat NaN poison
    # described in ``sample_start_times``.
    frame_selection = torch.floor(self._motion_buffer_start_time_s[assigned_ids] * data.framerate[motion_ids]).to(
        torch.long
    )
    frame_selection = torch.minimum(frame_selection, (data.buffer_length[motion_ids] - 1).to(torch.long))

    # Same shared flat-index gather as ``_copy_frame_data``: the reset
    # path pays the identical per-field ConcatBatchTensor overhead otherwise.
    fast_fields = _fast_gather_fields(data)
    flat_rows: torch.Tensor | None = None

    def _gather(src, field: str) -> torch.Tensor:
        nonlocal flat_rows
        if field in fast_fields:
            if flat_rows is None:
                flat_rows = src.batch_starts[motion_ids] + frame_selection
            return src.concatenated_tensor.index_select(0, flat_rows)
        return src[motion_ids, frame_selection]

    for field in (
        "base_quat_w",
        "base_lin_vel_w",
        "base_ang_vel_w",
        "dofs_pos",
        "dofs_vel",
    ):
        src = getattr(data, field, None)
        dst = getattr(state_buffer, field, None)
        if src is None or dst is None:
            continue
        dst[envs_idx] = _gather(src, field)

    if data.base_pos_w is not None and state_buffer.base_pos_w is not None:
        base_pos_w = _gather(data.base_pos_w, "base_pos_w")
        # NOTE: move the base position to the origin of the motion
        if env_origins is not None:
            base_pos_w = base_pos_w + env_origins[envs_idx]
        # TODO: check if any link is below the ground
        state_buffer.base_pos_w[envs_idx] = base_pos_w

compute

compute(motion_sequence_key: str, envs_idx: slice | Sequence[int] | Tensor, sample_timestamp: Tensor, env_origins: Tensor | None, data_buffer: MotionReferenceData) -> None

Fill the motion data into the buffer for the specified envs only.

This should not modify any other data in the MotionReferenceData.

PARAMETER	DESCRIPTION
motion_sequence_key	The motion sequence key to sample the motion data from. TYPE: str
envs_idx	The environment IDs to fill the motion data. A slice (the per-step fill of every env) is gathered straight into the buffer, with no intermediate copy. TYPE: slice | Sequence[int] | Tensor
sample_timestamp	The time (s) to sample from the start of each env, shape (num_envs_idx, num_frames). TYPE: Tensor
env_origins	The origins for each robot/env, which is also affecting the motion reference data, shape (num_envs, 3). None: every env's origin is the world origin. TYPE: Tensor | None
data_buffer	The MotionReferenceData to fill the motion data, which is a buffer for all envs, shape (num_envs, num_frames, ...). TYPE: MotionReferenceData

Source code in eden/managers/motion_manager.py

def compute(
    self,
    motion_sequence_key: str,
    envs_idx: slice | Sequence[int] | torch.Tensor,
    sample_timestamp: torch.Tensor,
    env_origins: torch.Tensor | None,
    data_buffer: MotionReferenceData,
) -> None:
    """Fill the motion data into the buffer for the specified envs only.

    This should not modify any other data in the MotionReferenceData.

    Parameters
    ----------
    motion_sequence_key : str
        The motion sequence key to sample the motion data from.
    envs_idx : slice | Sequence[int] | torch.Tensor
        The environment IDs to fill the motion data. A slice (the per-step fill of every env) is
        gathered straight into the buffer, with no intermediate copy.
    sample_timestamp : torch.Tensor
        The time (s) to sample from the start of each env, shape (num_envs_idx, num_frames).
    env_origins : torch.Tensor | None
        The origins for each robot/env, which is also affecting the motion reference data, shape (num_envs, 3).
        None: every env's origin is the world origin.
    data_buffer : MotionReferenceData
        The MotionReferenceData to fill the motion data, which is a buffer for all envs,
        shape (num_envs, num_frames, ...).
    """
    assigned_ids = self.envs_idx_to_assigned_ids(envs_idx)
    data = self.data_source.motion_sequences[motion_sequence_key]

    motion_ids, frame_idx, validity = self._select_frames(
        assigned_ids, sample_timestamp, data, data_buffer, envs_idx
    )
    origin_offset = None if env_origins is None else env_origins[envs_idx]
    self._copy_frame_data(data, data_buffer, envs_idx, motion_ids, frame_idx, validity, origin_offset)

assigned_ids_from_envs_idx

assigned_ids_from_envs_idx(envs_idx: slice | Sequence[int] | Tensor | None = None) -> Tensor

Local int ids in [0, num_assigned_envs) for an int-like env selector.

Accepts None (all assigned envs), a slice, or a long/int tensor (or Sequence[int]) of global env ids. A bool mask is not accepted here — its assigned-space form is a bool mask, not int ids; use :meth:assigned_mask_from_envs_mask, or :meth:envs_idx_to_assigned_ids when the selector dtype is only known at runtime.

Source code in eden/managers/motion_manager.py

def assigned_ids_from_envs_idx(
    self,
    envs_idx: slice | Sequence[int] | torch.Tensor | None = None,
) -> torch.Tensor:
    """Local int ids in ``[0, num_assigned_envs)`` for an *int-like* env selector.

    Accepts ``None`` (all assigned envs), a slice, or a long/int tensor (or
    ``Sequence[int]``) of **global** env ids. A bool mask is **not** accepted
    here — its assigned-space form is a bool mask, not int ids; use
    :meth:`assigned_mask_from_envs_mask`, or :meth:`envs_idx_to_assigned_ids`
    when the selector dtype is only known at runtime.
    """
    if envs_idx is None:
        return torch.arange(self.num_assigned_envs, device=self.device)
    if isinstance(envs_idx, slice):
        start = envs_idx.start if envs_idx.start is not None else self.assigned_env_slice.start
        stop = envs_idx.stop if envs_idx.stop is not None else self.assigned_env_slice.stop
        step = envs_idx.step if envs_idx.step is not None else 1
        indices = torch.arange(start, stop, step, device=self.device, dtype=torch.long)
        return indices - self.assigned_env_slice.start
    if not isinstance(envs_idx, torch.Tensor):
        envs_idx = torch.tensor(envs_idx, device=self.device, dtype=torch.long)
    result = envs_idx - self.assigned_env_slice.start
    # Reading the check back is a GPU->CPU sync, and this runs on every step's reset and fill.
    # A buffer that starts at env 0 can only go below it on a negative id, which no in-tree
    # caller makes; device ids are trusted there, as `sanitize_envs_idx` trusts them.
    if self.assigned_env_slice.start > 0 or result.device.type == "cpu":
        assert (result >= 0).all(), (
            f"envs_idx contains indices below assigned_env_slice.start ({self.assigned_env_slice.start})"
        )
    return result

assigned_mask_from_envs_mask

assigned_mask_from_envs_mask(envs_mask: Tensor) -> Tensor

Local bool mask over [0, num_assigned_envs) for a global bool mask.

The assigned-space form is the global mask sliced to the assigned env range, kept as a bool mask (not int ids) so downstream resampling can over-sample the full buffer and mask-write without a GPU→CPU sync.

Source code in eden/managers/motion_manager.py

def assigned_mask_from_envs_mask(self, envs_mask: torch.Tensor) -> torch.Tensor:
    """Local bool mask over ``[0, num_assigned_envs)`` for a *global* bool mask.

    The assigned-space form is the global mask sliced to the assigned env
    range, kept as a bool mask (not int ids) so downstream resampling can
    over-sample the full buffer and mask-write without a GPU→CPU sync.
    """
    return envs_mask[self.assigned_env_slice.start : self.assigned_env_slice.stop]

envs_idx_to_assigned_ids

envs_idx_to_assigned_ids(envs_idx: slice | Sequence[int] | Tensor | None = None) -> Tensor

Dispatch a runtime-typed env selector to its assigned-space form.

Returns a local bool mask when envs_idx is a bool tensor, otherwise local int ids. This dual return shape is load-bearing: a bool selector stays a bool mask so callers can mask-write (no GPU→CPU sync), while int/slice/None selectors become int ids that index local buffers directly. When the selector dtype is statically known, call :meth:assigned_ids_from_envs_idx or :meth:assigned_mask_from_envs_mask directly instead of going through this dispatcher.

Source code in eden/managers/motion_manager.py

def envs_idx_to_assigned_ids(
    self,
    envs_idx: slice | Sequence[int] | torch.Tensor | None = None,
) -> torch.Tensor:
    """Dispatch a *runtime-typed* env selector to its assigned-space form.

    Returns a **local bool mask** when ``envs_idx`` is a bool tensor, otherwise
    **local int ids**. This dual return shape is load-bearing: a bool selector
    stays a bool mask so callers can mask-write (no GPU→CPU sync), while
    int/slice/``None`` selectors become int ids that index local buffers
    directly. When the selector dtype is statically known, call
    :meth:`assigned_ids_from_envs_idx` or :meth:`assigned_mask_from_envs_mask`
    directly instead of going through this dispatcher.
    """
    if is_env_mask(envs_idx):
        return self.assigned_mask_from_envs_mask(envs_idx)
    return self.assigned_ids_from_envs_idx(envs_idx)

envs_idx_is_assigned

envs_idx_is_assigned(envs_idx: slice | Tensor) -> Tensor

Check if the envs_idx are assigned to this motion buffer.

PARAMETER	DESCRIPTION
envs_idx	The environment IDs to check. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	A boolean tensor indicating assignment.

Source code in eden/managers/motion_manager.py

def envs_idx_is_assigned(self, envs_idx: slice | torch.Tensor) -> torch.Tensor:
    """
    Check if the envs_idx are assigned to this motion buffer.

    Parameters
    ----------
    envs_idx : torch.Tensor
        The environment IDs to check.

    Returns
    -------
    torch.Tensor
        A boolean tensor indicating assignment.
    """
    return (envs_idx >= self.assigned_env_slice.start) & (envs_idx < self.assigned_env_slice.stop)

get_current_motion_weights

get_current_motion_weights(envs_idx: slice | Tensor | None = None) -> Tensor

Get the current motion weights for the specified envs.

The motion weights are used to determine the probability of selecting a motion reference for each env.

_assigned_env_motion_selection holds local motion ids; the canonical _motion_weights is global-sized, so we go through _local_to_global to index correctly.

PARAMETER	DESCRIPTION
envs_idx	The environment IDs (or bool mask) to get weights for; None for every assigned env. TYPE: slice | Tensor | None DEFAULT: None

Source code in eden/managers/motion_manager.py

def get_current_motion_weights(self, envs_idx: slice | torch.Tensor | None = None) -> torch.Tensor:
    """Get the current motion weights for the specified envs.

    The motion weights are used to determine the probability of selecting a motion reference for each env.

    ``_assigned_env_motion_selection`` holds local motion ids; the
    canonical ``_motion_weights`` is global-sized, so we go through
    ``_local_to_global`` to index correctly.

    Parameters
    ----------
    envs_idx : slice | torch.Tensor | None, optional
        The environment IDs (or bool mask) to get weights for; None for every assigned env.
    """
    if envs_idx is not None:
        assigned_ids = self.envs_idx_to_assigned_ids(envs_idx)
        local_ids = self._assigned_env_motion_selection[assigned_ids]
    else:
        local_ids = self._assigned_env_motion_selection
    return self._motion_weights[self._local_to_global[local_ids]]

match_scene

match_scene(entity_name: str)

Let the motion buffer match the scene.

For example a scene-interactive motion may need to match the motion starting point to the terrain. Recommended to call scene.terrain.terrain_origins and scene.terrain.subterrain_specific_cfgs to get some structured terrain information.

PARAMETER	DESCRIPTION
entity_name	The name of the entity whose motion should be matched to the scene. TYPE: str

Source code in eden/managers/motion_manager.py

def match_scene(self, entity_name: str):
    """Let the motion buffer match the scene.

    For example a scene-interactive motion may need to match the motion starting point to the terrain.
    Recommended to call `scene.terrain.terrain_origins` and `scene.terrain.subterrain_specific_cfgs`
    to get some structured terrain information.

    Parameters
    ----------
    entity_name : str
        The name of the entity whose motion should be matched to the scene.
    """
    pass

summarize_per_motion_metrics

summarize_per_motion_metrics() -> dict[str, Tensor]

Per-motion-id tensors for CSV dump; {} when metrics are disabled.

Public accessor for external dump tooling (see MotionTermOptions.compute_per_motion_metrics); delegates to the PerMotionMetrics collaborator that owns the buffers.

Source code in eden/managers/motion_manager.py

def summarize_per_motion_metrics(self) -> dict[str, torch.Tensor]:
    """Per-motion-id tensors for CSV dump; ``{}`` when metrics are disabled.

    Public accessor for external dump tooling (see
    ``MotionTermOptions.compute_per_motion_metrics``); delegates to the
    ``PerMotionMetrics`` collaborator that owns the buffers.
    """
    return self._metrics.per_motion_tensors() if self._metrics is not None else {}

MotionManager

MotionManager(env: EnvBase, options: MotionManagerOptions)

Bases: FlatManagerBase[MotionManagerOptions], SymmetricAugmentationMixin

The manager to handle all motion references.

Source code in eden/managers/motion_manager.py

def __init__(
    self,
    env: EnvBase,
    options: MotionManagerOptions,
):
    # NOTE: these must be set BEFORE super().__init__() because
    # FlatManagerBase.__init__ -> ManagerBase.__init__ calls _prepare_terms() which uses them.
    self._is_visualizing = False
    self._timestamp = torch.zeros(env.num_envs, device=env.device)
    # Opt-in motion-end resync (keeps episodes running past a clip's end).
    # No-op unless options.resync_on_motion_end; owns its own state.
    self._resync = MotionEndResync(self)
    self._frame_interval_s = torch.ones(env.num_envs, device=env.device) * options.frame_interval_s_range[0]
    self.require_frame_interval_s_resample = (
        options.frame_interval_s_range[1] - options.frame_interval_s_range[0]
    ) > 0

    self._env_symmetric_augmentation_mask = None
    if options.symmetric_augmentation:
        grouped = [
            name
            for name in options.term_keys()
            if getattr(getattr(options, name), "env_motion_groups", None) is not None
        ]
        if grouped:
            # The mirror applies one joint/link map per entity: it cannot swap a left and a right hand that are
            # separate entities, nor mirror the object a grouped clip manipulates.
            raise ValueError(
                f"symmetric_augmentation=True cannot be combined with env_motion_groups (set on motion term(s) "
                f"{grouped}): mirroring a clip would not mirror the per-variant entity it is grouped by. Set "
                f"motion_options.symmetric_augmentation=False."
            )
        self._env_symmetric_augmentation_mask = torch.zeros(env.num_envs, device=env.device, dtype=torch.bool)
    # The envs of the reset in progress whose init state is frame 0 of their window, and whether that
    # window has been filled yet (see ``resetting``). None outside a reset.
    self._resetting_envs: torch.Tensor | None = None
    self._resetting_window_filled = False
    # Inside ``managers_resetting``: the manager resets of a reset, during which ``data`` raises.
    self._managers_resetting = False
    # Per entity a reset-time reader asked for (see ``record_ended_episode_validity``): the window's
    # ``validity[:, 0]`` as each env's last episode ended.
    self._ended_episode_validity: dict[str, torch.Tensor] = {}

    super().__init__(env=env, options=options)

    # Fail fast: when motion-end resync is enabled, resolve + validate its RSI teleport
    # terms via a post-build hook (the event manager is built after this manager), so a
    # missing RSI term raises at env-construction time rather than at the first wrap.
    if self._resync.enabled:
        env.register_post_build_hook(self._resync.validate)

Attributes

require_frame_interval_s_resample instance-attribute

require_frame_interval_s_resample = options.frame_interval_s_range[1] - options.frame_interval_s_range[0] > 0

data property

data: dict[str, MotionReferenceData]

The motion reference data for all entities, keyed by entity name.

Filled by :meth:compute at the start of each MimicRLEnvBase.step() and refilled per-env by :meth:resetting in :meth:MimicRLEnvBase._reset_idx. Fresh for all envs during reward, termination, action, metric, and observation computation in RLEnvBase.step().

Stale window: :meth:reset is called inside MimicRLEnvBase._reset_idx before super()._reset_idx and only resamples motion selection / rolls back _timestamp — it does not refill the buffer. :meth:resetting, which brackets the scene reset and the RESET events, does: at the first init-state request inside it (data_start_from="current_time"), else on its exit. A RESET-mode event may therefore see stale data for reset envs (last step's compute fill); use :meth:get_reset_reference_state instead — it is the one ReferenceStateInitialization uses.

Raises inside the manager resets (:meth:managers_resetting): the reset envs' window already holds the new episode there, so a reset() reading it to grade the ended episode would grade the wrong one, silently. Grade the ended episode off :meth:ended_episode_validity; read the new episode's reference in post_reset().

reference_frame property

reference_frame: dict[str, MotionReferenceData]

The reference state of all entities at current time step, keyed by entity name.

num_frames property

num_frames: int

env_origins property

env_origins: Tensor | None

The env's env_origins (num_envs, 3), which shift the reference; None for an env without them.

None, not zeros, so that a fill does not add zeros to every position field of every entity.

has_adaptive_motion_weights property

has_adaptive_motion_weights: bool

True iff at least one motion term enables hard-motion mining.

Methods:

summary

summary()

Get the tabular information of the motion reference managed buffer.

Source code in eden/managers/motion_manager.py

def summary(self):
    """Get the tabular information of the motion reference managed buffer."""
    msg = f"<MotionManager> contains {len(self._terms)} motion terms to use.\n"

    msg += format_table(
        "Active Motion Terms",
        ["Name", "Num Trajectories"],
        ([name, term.num_trajectories] for name, term in self._terms.items()),
        align={"Name": "l", "Num Trajectories": "r"},
    )
    msg += "\n"

    return msg

reset

reset(envs_idx: Sequence[int] | Tensor | None = None) -> dict[str, Any]

Reset the motion reference manager.

Resamples motion selection, rolls back _timestamp[envs_idx] to 0, and updates adaptive-sampling bookkeeping. Does not refill self._data — callers that need fresh per-step data must call :meth:fill_data themselves. MimicRLEnvBase does it through :meth:resetting, around the scene reset and RESET events (so the refill picks up any env_origins update from scene.reset / RESET events, as the reference-state seeding does).

Source code in eden/managers/motion_manager.py

def reset(self, envs_idx: Sequence[int] | torch.Tensor | None = None) -> dict[str, Any]:
    """Reset the motion reference manager.

    Resamples motion selection, rolls back ``_timestamp[envs_idx]`` to 0,
    and updates adaptive-sampling bookkeeping. **Does not refill**
    ``self._data`` — callers that need fresh per-step data must call
    :meth:`fill_data` themselves. ``MimicRLEnvBase`` does it through
    :meth:`resetting`, around the scene reset and RESET events (so the
    refill picks up any ``env_origins`` update from ``scene.reset`` /
    RESET events, as the reference-state seeding does).
    """
    if envs_idx is None:
        envs_idx = self._all_indices
    else:
        envs_idx = torch.as_tensor(envs_idx, device=self.device)

    self._record_adaptive_bin_outcomes(envs_idx)
    self._record_motion_outcomes(envs_idx)
    for name, ended in self._ended_episode_validity.items():
        assign_over_envs(ended, envs_idx, self._data[name].validity[:, 0])
    self._timestamp[envs_idx] = 0.0
    self._reset_motion_terms(envs_idx)
    self._resample_buffer_collate_params(envs_idx)

    return self._collect_log_extras()

post_reset

post_reset(envs_idx: slice | Tensor | None = None) -> None

Fan out to motion terms, restricted to each term's assigned environments.

Overrides the flat fan-out (which would hand every term the full reset index) because a motion term owns a contiguous env slice and its buffers are sized to that slice — the same per-term narrowing :meth:compute and :meth:fill_data do. Terms are unique per name here, so the base class's dedupe is not needed.

Source code in eden/managers/motion_manager.py

@torch.inference_mode()
def post_reset(self, envs_idx: slice | torch.Tensor | None = None) -> None:
    """Fan out to motion terms, restricted to each term's assigned environments.

    Overrides the flat fan-out (which would hand every term the full reset index)
    because a motion term owns a contiguous env slice and its buffers are sized to
    that slice — the same per-term narrowing :meth:`compute` and :meth:`fill_data`
    do. Terms are unique per name here, so the base class's dedupe is not needed.
    """
    envs_idx = sanitize_envs_idx(envs_idx, self.num_envs, prefer_slice=False)
    for name, term in self._terms.items():
        envs_idx_this_term, _ = self._filter_envs_for_term(envs_idx, name)
        term.post_reset(envs_idx_this_term)

collect_per_motion_distributions

collect_per_motion_distributions(wandb_module) -> dict[str, 'Any']

Per-term per-motion histograms for the wandb-direct emit path.

Kept off the _collect_log_extras / extras["log"] channel: those flow through rsl_rl's Logger.ep_extras which iterates and prints every key per iter — at AMASS scale that loop is the per-step bottleneck. The wandb-direct mixin path (see MotionWeightCurriculumMixin) calls this method once per learning iter and merges the result into the same wandb.log() payload. Empty when no term has log_per_motion_metrics=True.

Source code in eden/managers/motion_manager.py

def collect_per_motion_distributions(self, wandb_module) -> dict[str, "Any"]:
    """Per-term per-motion histograms for the wandb-direct emit path.

    Kept off the ``_collect_log_extras`` / ``extras["log"]`` channel: those
    flow through rsl_rl's ``Logger.ep_extras`` which iterates and prints
    every key per iter — at AMASS scale that loop is the per-step
    bottleneck. The wandb-direct mixin path (see
    ``MotionWeightCurriculumMixin``) calls this method once per learning
    iter and merges the result into the same ``wandb.log()`` payload.
    Empty when no term has ``log_per_motion_metrics=True``.
    """
    out: dict[str, "Any"] = {}
    if not self._terms:
        return out
    multi_term = len(self._terms) > 1
    for name, term in self._terms.items():
        if term._curriculum is None:
            continue
        sub = term._curriculum.collect_per_motion_distributions(wandb_module)
        if not sub:
            continue
        if multi_term:
            sub = {f"{name}/{k}": v for k, v in sub.items()}
        out.update(sub)
    return out

compute

compute(dt: float)

Advance _timestamp and refill self._data for all envs.

Called once per MimicRLEnvBase.step (before reward / termination / observation compute). The reset path fills the reset-env subset through :meth:resetting (see :meth:fill_data).

When any motion term sets freeze_at_timestep_zero_prob > 0, envs currently at _timestamp == 0 and assigned to that term are per-step-randomly held at 0 (re-rolled every step — see :class:MotionTermOptions.freeze_at_timestep_zero_prob).

Source code in eden/managers/motion_manager.py

def compute(self, dt: float):
    """Advance ``_timestamp`` and refill ``self._data`` for all envs.

    Called once per ``MimicRLEnvBase.step`` (before reward / termination
    / observation compute). The reset path fills the reset-env subset
    through :meth:`resetting` (see :meth:`fill_data`).

    When any motion term sets ``freeze_at_timestep_zero_prob > 0``, envs
    currently at ``_timestamp == 0`` and assigned to that term are
    per-step-randomly held at 0 (re-rolled every step — see
    :class:`MotionTermOptions.freeze_at_timestep_zero_prob`).
    """
    self._advance_timestamp(dt)
    # Blend the death-bin failure counts recorded at the previous reset
    # (see ``AdaptiveBinSampler.record_bin_outcomes``) into their EMA.
    self._blend_adaptive_ema()
    # Motion-end resync (opt-in), folded into the single per-step fill_data(all):
    # pre_fill detects wrapped envs from the advanced clock (before fill) and
    # resamples them so the ONE fill_data below fills them at the new phase (no
    # second fill_data); post_fill then teleports them onto it, through the wrap
    # mask, so no step reads on the host whether any env wrapped. None
    # when disabled. See eden.managers.motion.resync.MotionEndResync.
    wrap_mask = self._resync.pre_fill()
    self.fill_data()
    if wrap_mask is not None:
        self._resync.post_fill(wrap_mask)
    self._accumulate_per_motion_metrics()

rewind_to_start

rewind_to_start(envs_idx: Tensor) -> None

Rewind playback to frame 0 for envs_idx WITHOUT resampling.

Distinct from :meth:reset: reset re-samples motion selection and start phase (the per-episode randomized restart), whereas this replays each env's CURRENTLY-assigned clip deterministically from the beginning (e.g. a deploy starting from a stand). Zeros the manager timestamp and fans out to each term's :meth:MotionTerm.rewind_to_start. Does not refill the reference buffer — call :meth:fill_data afterwards.

Source code in eden/managers/motion_manager.py

def rewind_to_start(self, envs_idx: torch.Tensor) -> None:
    """Rewind playback to frame 0 for ``envs_idx`` WITHOUT resampling.

    Distinct from :meth:`reset`: ``reset`` re-samples motion selection and start phase
    (the per-episode randomized restart), whereas this replays each env's CURRENTLY-assigned
    clip deterministically from the beginning (e.g. a deploy starting from a stand). Zeros
    the manager timestamp and fans out to each term's :meth:`MotionTerm.rewind_to_start`.
    Does not refill the reference buffer — call :meth:`fill_data` afterwards.
    """
    self._timestamp[envs_idx] = 0.0
    # MOTION_TERM_REGISTRY enforces MotionTerm, so every term has rewind_to_start.
    for term in self._terms.values():
        term.rewind_to_start(envs_idx)

fill_data

fill_data(envs_idx: Tensor | None = None) -> None

Refill self._data[entity][envs_idx] from self._timestamp[envs_idx]; None fills every env.

Called by :meth:compute (per step, every env) and by :meth:resetting (reset envs only), around the scene reset and RESET events of MimicRLEnvBase._reset_idx. Bool-mask input becomes int ids through :func:~eden.utils.misc.env_indices — free for the auto-reset's mask, whose ids the env took once for the whole reset.

Only the filled envs' rows are written. A fill of every env (None) hands each term its assignment as a slice, which the term gathers straight into the buffer; the ids of every env write the same rows, through the ids.

Source code in eden/managers/motion_manager.py

def fill_data(self, envs_idx: torch.Tensor | None = None) -> None:
    """Refill ``self._data[entity][envs_idx]`` from ``self._timestamp[envs_idx]``; None fills every env.

    Called by :meth:`compute` (per step, every env) and by :meth:`resetting` (reset envs only),
    around the scene reset and RESET events of ``MimicRLEnvBase._reset_idx``. Bool-mask input
    becomes int ids through :func:`~eden.utils.misc.env_indices` — free for the auto-reset's mask,
    whose ids the env took once for the whole reset.

    Only the filled envs' rows are written. A fill of every env (None) hands each term its
    assignment as a slice, which the term gathers straight into the buffer; the ids of
    every env write the same rows, through the ids.
    """
    every_env = envs_idx is None
    envs_idx = self._all_indices if every_env else env_indices(envs_idx)

    n = envs_idx.shape[0]
    if n == 0:
        return

    time_to_target_frame = torch.arange(self.num_frames, device=self.device, dtype=torch.float32)
    if self._options.data_start_from == "one_frame_interval":
        time_to_target_frame += 1
    time_to_target_frame = time_to_target_frame.unsqueeze(0).repeat(n, 1)
    time_to_target_frame *= self._frame_interval_s[envs_idx].unsqueeze(-1)

    for name, term in self._terms.items():
        if every_env:
            # Row i of every env's tensors is env i, so the term's envs are its assignment slice.
            envs_idx_this_buffer = self._motion_term_assignment[name]
            time_to_target_frame_this_buffer = time_to_target_frame[envs_idx_this_buffer]
            ids_this_buffer = self._all_indices[envs_idx_this_buffer]
        else:
            envs_idx_this_buffer, envs_idx_this_buffer_mask = self._filter_envs_for_term(envs_idx, name)
            time_to_target_frame_this_buffer = (
                time_to_target_frame
                if envs_idx_this_buffer_mask is None
                else time_to_target_frame[envs_idx_this_buffer_mask]
            )
            ids_this_buffer = envs_idx_this_buffer
        for entity_name, motion_seq_key in term.entities_map.items():
            term.compute(
                motion_seq_key,
                envs_idx_this_buffer,
                self._timestamp[envs_idx_this_buffer].unsqueeze(-1) + time_to_target_frame_this_buffer,
                self.env_origins,
                self._data[entity_name],
            )
            if self._symmetric_augmentation_configured:
                self._symmetric_augment_reference_data(self._data[entity_name], ids_this_buffer)

    rows = slice(None) if every_env else envs_idx
    for entity_name, entity_data in self._data.items():
        # An invalid frame's time is -1. The other envs' rows are as their own fill left them.
        entity_data.time_to_target_frame[rows] = torch.where(entity_data.validity[rows], time_to_target_frame, -1.0)

get_target_dofs_name

get_target_dofs_name(entity_name: str) -> list[str] | None

Return the resolved DOF names that the motion reference exposes for entity_name.

These are the names that index into reference_frame[entity_name].dofs_pos along the last axis. Returns None when no motion term targets the entity.

When multiple terms target the same entity, the first term is used; the :class:MotionManager already requires consistent num_dofs across terms.

Source code in eden/managers/motion_manager.py

def get_target_dofs_name(self, entity_name: str) -> list[str] | None:
    """Return the resolved DOF names that the motion reference exposes for ``entity_name``.

    These are the names that index into ``reference_frame[entity_name].dofs_pos``
    along the last axis. Returns ``None`` when no motion term targets the entity.

    When multiple terms target the same entity, the first term is used; the
    :class:`MotionManager` already requires consistent num_dofs across terms.
    """
    for term in self._terms.values():
        motion_seq_key = term.entities_map.get(entity_name)
        if motion_seq_key is None:
            continue
        names = term.target_dofs_name.get(motion_seq_key)
        if isinstance(names, str):
            return [names]
        return list(names) if names is not None else []
    return None

get_motion_phase

get_motion_phase(entity_name: str) -> Tensor

Per-env normalized playback phase of the clip tracking entity_name.

frame / (clip_length - 1) in [0, 1] at the CURRENT reference frame (the same floor-of-timestamp frame selection as :meth:MotionTerm._select_frames), so mid-clip RSI starts report their true position in the clip rather than the episode-step phase. Envs not assigned to a term tracking the entity report 0. Shape (num_envs,).

Raises if no term tracks entity_name at all: that is a config error (a misspelled entity), and returning the all-zero buffer would hand a consumer a silently-constant signal forever.

Source code in eden/managers/motion_manager.py

def get_motion_phase(self, entity_name: str) -> torch.Tensor:
    """Per-env normalized playback phase of the clip tracking ``entity_name``.

    ``frame / (clip_length - 1)`` in ``[0, 1]`` at the CURRENT reference frame
    (the same floor-of-timestamp frame selection as :meth:`MotionTerm._select_frames`),
    so mid-clip RSI starts report their true position in the clip rather than
    the episode-step phase. Envs not assigned to a term tracking the entity
    report 0. Shape ``(num_envs,)``.

    Raises if *no* term tracks ``entity_name`` at all: that is a config error
    (a misspelled entity), and returning the all-zero buffer would hand a
    consumer a silently-constant signal forever.
    """
    tracking_terms = [term for term in self._terms.values() if entity_name in term.entities_map]
    if not tracking_terms:
        tracked_names = sorted({name for term in self._terms.values() for name in term.entities_map})
        raise ValueError(
            f"Entity '{entity_name}' is not tracked by any motion term, so its playback phase is "
            f"undefined. Tracked entities: {tracked_names}"
        )
    phase = torch.zeros(self.num_envs, device=self.device)
    for term in tracking_terms:
        data = term.data_source.motion_sequences[term.entities_map[entity_name]]
        assigned = term.assigned_env_slice
        selection = term._assigned_env_motion_selection
        length = data.buffer_length[selection].to(torch.float32)
        frame = torch.floor(
            (term._motion_buffer_start_time_s + self._timestamp[assigned]) * data.framerate[selection]
        )
        phase[assigned] = torch.clamp(frame, max=length - 1.0) / (length - 1.0).clamp(min=1.0)
    return phase

get_init_reference_state

get_init_reference_state(envs_idx: Sequence[int] | Tensor | None = None, entity_names: Sequence[str] | None = None) -> dict[str, MotionReferenceState]

Get the initial reference state for the given envs_idx, keyed by entity name.

entity_names limits the gather to those entities (None: every tracked one). Each entity costs one gather per field, so the reset path asks only for the entities it seeds rather than refilling all of them once per RSI term. Valid for any envs, at any time; a RESET event seeding its own envs_idx calls :meth:get_reset_reference_state, which saves the gather.

Source code in eden/managers/motion_manager.py

def get_init_reference_state(
    self,
    envs_idx: Sequence[int] | torch.Tensor | None = None,
    entity_names: Sequence[str] | None = None,
) -> dict[str, MotionReferenceState]:
    """Get the initial reference state for the given envs_idx, keyed by entity name.

    ``entity_names`` limits the gather to those entities (``None``: every tracked one). Each
    entity costs one gather per field, so the reset path asks only for the entities it seeds
    rather than refilling all of them once per RSI term. Valid for any envs, at any time;
    a RESET event seeding its own ``envs_idx`` calls :meth:`get_reset_reference_state`, which
    saves the gather.
    """
    wanted = self._wanted_entities(entity_names, "get_init_reference_state")
    if envs_idx is None:
        envs_idx = self._all_indices
    else:
        # Int ids, not a bool mask: every field below is gathered and written for the selected
        # envs, a GPU->CPU sync per op through a mask and none through ids.
        envs_idx = env_indices(torch.as_tensor(envs_idx, device=self.device))
    for name, term in self._terms.items():
        entities = [(entity, key) for entity, key in term.entities_map.items() if entity in wanted]
        if not entities:
            continue
        envs_idx_this_buffer, _ = self._filter_envs_for_term(envs_idx, name)
        for entity_name, motion_seq_key in entities:
            term.fill_init_reference_state(
                motion_seq_key,
                envs_idx_this_buffer,
                self.env_origins,
                self._init_reference_state[entity_name],
            )
            if self._symmetric_augmentation_configured:
                self._symmetric_augment_reference_data(
                    self._init_reference_state[entity_name],
                    envs_idx_this_buffer,
                )
    return {name: state[envs_idx] for name, state in self._init_reference_state.items() if name in wanted}

get_reset_reference_state

get_reset_reference_state(envs_idx: Sequence[int] | Tensor | None = None, entity_names: Sequence[str] | None = None) -> dict[str, MotionReferenceState]

:meth:get_init_reference_state for a RESET event seeding its own envs_idx.

Inside :meth:resetting with data_start_from="current_time", the window's first frame is each reset env's init state, so the first call fills the reset envs' window and every call reads the state off its frame 0: the reset gathers each reference row once, not once for the window and again for reference-state initialization. Anywhere else it is :meth:get_init_reference_state.

envs_idx must be envs of the reset in progress, which a RESET event's envs_idx always is (a throttled event's is a subset). Frame 0 of any other env's window is a mid-episode reference row, and it is not checked here, since that would cost a GPU->CPU sync every reset: call :meth:get_init_reference_state for envs that need not be resetting.

Source code in eden/managers/motion_manager.py

def get_reset_reference_state(
    self,
    envs_idx: Sequence[int] | torch.Tensor | None = None,
    entity_names: Sequence[str] | None = None,
) -> dict[str, MotionReferenceState]:
    """:meth:`get_init_reference_state` for a RESET event seeding its own ``envs_idx``.

    Inside :meth:`resetting` with ``data_start_from="current_time"``, the window's first frame *is*
    each reset env's init state, so the first call fills the reset envs' window and every call
    reads the state off its frame 0: the reset gathers each reference row once, not once for the
    window and again for reference-state initialization. Anywhere else it is
    :meth:`get_init_reference_state`.

    ``envs_idx`` must be envs of the reset in progress, which a RESET event's ``envs_idx`` always
    is (a throttled event's is a subset). Frame 0 of any other env's window is a mid-episode
    reference row, and it is not checked here, since that would cost a GPU->CPU sync every reset:
    call :meth:`get_init_reference_state` for envs that need not be resetting.
    """
    if self._resetting_envs is None or envs_idx is None:
        return self.get_init_reference_state(envs_idx, entity_names)
    wanted = self._wanted_entities(entity_names, "get_reset_reference_state")
    # Int ids, as in ``get_init_reference_state``.
    return self._init_state_from_window(env_indices(torch.as_tensor(envs_idx, device=self.device)), wanted)

record_ended_episode_validity

record_ended_episode_validity(entity_name: str) -> None

Keep, from now on, entity_name's :meth:ended_episode_validity (call it in a term's build()).

For a term that grades the episode a reset ends: the manager resets run after the reset envs' window has been refilled for the new episode (see :meth:resetting), so they cannot read the ended one off data.

Source code in eden/managers/motion_manager.py

def record_ended_episode_validity(self, entity_name: str) -> None:
    """Keep, from now on, ``entity_name``'s :meth:`ended_episode_validity` (call it in a term's ``build()``).

    For a term that grades the episode a reset ends: the manager resets run after the reset envs'
    window has been refilled for the new episode (see :meth:`resetting`), so they cannot read the
    ended one off ``data``.
    """
    if entity_name not in self._data:
        raise ValueError(
            f"record_ended_episode_validity: no motion term tracks '{entity_name}'; tracked entities: "
            f"{sorted(self._data)}. Name one of them, or add the entity to a motion term's entities_map."
        )
    self._ended_episode_validity.setdefault(
        entity_name, torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
    )

ended_episode_validity

ended_episode_validity(entity_name: str) -> Tensor

validity[:, 0] of entity_name's window as each env's last episode ended, (num_envs,).

False where that episode played its clip out. Taken by :meth:reset, before the window moves on; an env that never ran an episode reads False, the empty window's validity.

Source code in eden/managers/motion_manager.py

def ended_episode_validity(self, entity_name: str) -> torch.Tensor:
    """``validity[:, 0]`` of ``entity_name``'s window as each env's last episode ended, ``(num_envs,)``.

    False where that episode played its clip out. Taken by :meth:`reset`, before the window moves
    on; an env that never ran an episode reads False, the empty window's validity.
    """
    ended = self._ended_episode_validity.get(entity_name)
    if ended is None:
        raise ValueError(
            f"ended_episode_validity('{entity_name}'): nothing recorded it. Call "
            f"motion_manager.record_ended_episode_validity('{entity_name}') in the term's build(); reading "
            f"data['{entity_name}'].validity in a reset() instead sees the new episode's window."
        )
    return ended

resetting

resetting(envs_idx: Tensor) -> Iterator[None]

Bracket the scene reset and RESET events of a reset of envs_idx, and leave its window filled.

Enter after :meth:reset; on a normal exit data holds the reset envs' new window, filled by :meth:fill_data once, so every manager reset after it sees the new episode (the ended one's clip state is :meth:ended_episode_validity). With data_start_from="current_time" the window's first frame is each reset env's init state, so the first :meth:get_reset_reference_state call inside fills the window there and every call reads the state off it. Otherwise, and when nothing asks, the window is filled on exit, after the RESET events. The window reads env_origins when it is filled: at the first such call, or on exit.

Source code in eden/managers/motion_manager.py

@contextmanager
def resetting(self, envs_idx: torch.Tensor) -> Iterator[None]:
    """Bracket the scene reset and RESET events of a reset of ``envs_idx``, and leave its window filled.

    Enter after :meth:`reset`; on a normal exit ``data`` holds the reset envs' new window, filled
    by :meth:`fill_data` once, so every manager reset after it sees the new episode (the ended one's
    clip state is :meth:`ended_episode_validity`). With ``data_start_from="current_time"`` the
    window's first frame *is* each reset env's init state, so the first
    :meth:`get_reset_reference_state` call inside fills the window there and every call reads the
    state off it. Otherwise, and when nothing asks, the window is filled on exit, after the
    RESET events. The window reads ``env_origins`` when it is filled: at the first such call, or
    on exit.
    """
    if self._options.data_start_from == "current_time":
        self._resetting_envs = envs_idx
    try:
        yield
        filled = self._resetting_window_filled
    finally:
        self._resetting_envs = None
        self._resetting_window_filled = False
    if not filled:
        self.fill_data(envs_idx)

managers_resetting

managers_resetting() -> Iterator[None]

Bracket the manager resets that follow :meth:resetting: reading :attr:data inside raises.

The reset envs' window holds the new episode by then, so a manager reset() reading it to grade the episode that ended would silently grade the new one.

Source code in eden/managers/motion_manager.py

@contextmanager
def managers_resetting(self) -> Iterator[None]:
    """Bracket the manager resets that follow :meth:`resetting`: reading :attr:`data` inside raises.

    The reset envs' window holds the new episode by then, so a manager ``reset()`` reading it to
    grade the episode that ended would silently grade the new one.
    """
    self._managers_resetting = True
    try:
        yield
    finally:
        self._managers_resetting = False

update_motion_weight_curriculum

update_motion_weight_curriculum() -> None

Run one curriculum step across all adaptive motion terms.

Called once per PPO iteration by the rsl-rl runner mixin (see eden.envs.wrappers.rsl_rl_env). Terms that haven't yet reached their per-term motion_weight_update_every_n_iters cadence are cheap no-ops; terms that have reached cadence drain their counts, all-reduce them across DDP ranks (so every rank applies the same update — a hard requirement for keeping _motion_weights byte-identical), then rewrite _motion_weights per motion_weight_update_mode.

DDP shape symmetry holds because every rank's outcome buffers and _motion_weights are sized to the GLOBAL motion-id space, independent of mp_split_method or each rank's local trajectory slice. Per-rank contributions land in the same global slots; after the SUM all_reduce every rank applies the same update and the resulting _motion_weights is byte-identical across ranks.

Source code in eden/managers/motion_manager.py

def update_motion_weight_curriculum(self) -> None:
    """Run one curriculum step across all adaptive motion terms.

    Called once per PPO iteration by the rsl-rl runner mixin (see
    ``eden.envs.wrappers.rsl_rl_env``). Terms that haven't yet reached
    their per-term ``motion_weight_update_every_n_iters`` cadence are
    cheap no-ops; terms that *have* reached cadence drain their counts,
    all-reduce them across DDP ranks (so every rank applies the same
    update — a hard requirement for keeping ``_motion_weights``
    byte-identical), then rewrite ``_motion_weights`` per
    ``motion_weight_update_mode``.

    DDP shape symmetry holds because every rank's outcome buffers and
    ``_motion_weights`` are sized to the GLOBAL motion-id space,
    independent of ``mp_split_method`` or each rank's local trajectory
    slice. Per-rank contributions land in the same global slots; after
    the SUM all_reduce every rank applies the same update and
    the resulting ``_motion_weights`` is byte-identical across ranks.
    """
    do_reduce = dist.is_initialized() and dist.get_world_size() > 1
    all_reduce_fn = (lambda t: dist.all_reduce(t, op=dist.ReduceOp.SUM)) if do_reduce else None
    for term in self._terms.values():
        term._update_motion_weight_curriculum(all_reduce_fn=all_reduce_fn)

Functions:

sample_start_times

sample_start_times(motion_ids: Tensor, motion_bin_weights: ConcatBatchTensor, buffer_device: device, motion_bin_length_s: float | None, data_source: 'DatasetBase | None' = None, uniform_bins: bool = False) -> Tensor

Sample a start time (seconds) per env from the bin-weight distribution.

uniform_bins=True is a fast path for when the bin weights are known to be uniform (i.e. adaptive_sampling is off, so the torch.ones init is never rewritten): a multinomial over uniform weights is exactly a uniform draw over each motion's valid bins, so we draw floor(rand * n_bins) directly and skip the padded get_with_mask gather + multinomial launch. This is bit-for-bit the same DISTRIBUTION as the weighted path (verified in tests/managers/test_motion_manager.py); the weighted path would dominate the motion-end resync's per-step cost, which calls this on every clip wrap.

When data_source is provided, results are clamped per-env to (buffer_length[motion_id] - 1) / framerate[motion_id]. This is the correct behavior whenever any clip is shorter than motion_bin_length_s — without the clamp, the (after-clamp(min=1)) single bin spans the full motion_bin_length_s window, so the intra-bin uniform noise can return a start time past the end of the clip. fill_init_reference_state then reads base_quat_w at an out-of-range frame index, ConcatBatchTensor zero-fills it to (0,0,0,0), and the first genesis quaternion op on that degenerate quat — the rotation math normalizes by q·q, so a zero quat is a divide-by-zero — NaN-poisons the policy observation on the first reset.

The sibling _sample_as_concat_motion_bins path applies the same clamp inline (see max_start_time_s below). Keeping data_source optional preserves the existing call signature for tests that build a CBT directly without a full DatasetBase; production callers should pass it.

Source code in eden/managers/motion_manager.py

def sample_start_times(
    motion_ids: torch.Tensor,
    motion_bin_weights: ConcatBatchTensor,
    buffer_device: torch.device,
    motion_bin_length_s: float | None,
    data_source: "DatasetBase | None" = None,
    uniform_bins: bool = False,
) -> torch.Tensor:
    """Sample a start time (seconds) per env from the bin-weight distribution.

    ``uniform_bins=True`` is a fast path for when the bin weights are known to be
    uniform (i.e. ``adaptive_sampling`` is off, so the ``torch.ones`` init is never
    rewritten): a ``multinomial`` over uniform weights is *exactly* a uniform draw
    over each motion's valid bins, so we draw ``floor(rand * n_bins)`` directly and
    skip the padded ``get_with_mask`` gather + ``multinomial`` launch. This is
    bit-for-bit the same DISTRIBUTION as the weighted path (verified in
    ``tests/managers/test_motion_manager.py``); the weighted path would dominate
    the motion-end resync's per-step cost, which calls this on every clip wrap.

    When ``data_source`` is provided, results are clamped per-env to
    ``(buffer_length[motion_id] - 1) / framerate[motion_id]``. This is the
    correct behavior whenever any clip is shorter than ``motion_bin_length_s``
    — without the clamp, the (after-``clamp(min=1)``) single bin spans the
    full ``motion_bin_length_s`` window, so the intra-bin uniform noise can
    return a start time past the end of the clip. ``fill_init_reference_state``
    then reads ``base_quat_w`` at an out-of-range frame index, ``ConcatBatchTensor``
    zero-fills it to ``(0,0,0,0)``, and the first genesis quaternion op on that
    degenerate quat — the rotation math normalizes by ``q·q``, so a zero quat is a
    divide-by-zero — NaN-poisons the policy observation on the first reset.

    The sibling ``_sample_as_concat_motion_bins`` path applies the same clamp
    inline (see ``max_start_time_s`` below). Keeping ``data_source`` optional
    preserves the existing call signature for tests that build a CBT directly
    without a full ``DatasetBase``; production callers should pass it.
    """
    if motion_bin_length_s is None:
        raise ValueError("motion_bin_length_s must be provided")

    num_assign = len(motion_ids)
    if num_assign == 0:
        return torch.empty(0, dtype=torch.float32, device=buffer_device)

    if uniform_bins:
        # Uniform weights => sampled bin is uniform over each motion's [0, n_bins).
        # batch_sizes[motion_ids] is that per-env bin count; floor(rand * n) draws it
        # directly (== multinomial over uniform weights) without the CBT gather.
        n_bins = motion_bin_weights.batch_sizes[motion_ids].to(torch.float)
        sampled_bins = torch.floor(torch.rand(num_assign, device=buffer_device) * n_bins)
    else:
        # Batched padded access: (num_assign, max_bins) with zeros for out-of-range bins
        # Padded to the longest motion's bin count, known on the host: the longest *selected* one is a sync.
        padded_weights, valid_mask = motion_bin_weights.get_with_mask(motion_ids, pad_to_longest_batch=True)
        # Guard against all-zero rows (would crash multinomial)
        safe_weights = padded_weights + valid_mask.float() * 1e-8

        # Single batched multinomial: one sample per env
        sampled_bins = torch.multinomial(safe_weights, 1, replacement=True).squeeze(-1).float()

    # Intra-bin uniform noise + scale to seconds
    start_times = (sampled_bins + torch.rand(num_assign, device=buffer_device)) * motion_bin_length_s
    if data_source is not None:
        max_start_time_s = (data_source.buffer_length[motion_ids].to(torch.float) - 1) / data_source.framerate[
            motion_ids
        ].to(torch.float)
        start_times = torch.clamp(start_times, max=max_start_time_s)
    return start_times

MotionTerm collaborators

The pieces MotionTerm delegates to: adaptive start-time bins, the per-motion outcome curriculum, per-motion tracking-error metrics, and the opt-in resync that keeps an episode running past a clip's end.

Intra-motion start-time bin curriculum (adaptive start-time sampling).

Split out of MotionTerm. Owns the per-motion temporal bin weights (a ConcatBatchTensor of per-bin sampling probabilities) and, when adaptive_sampling is on, the death-bin failure-count EMA + the vectorised weight-update that biases start-time sampling toward the run-up to hard moments.

One instance per term that sets motion_bin_length_s (created/recreated in MotionTerm.enable_trajectories whenever the motion library changes). The bin weights exist whenever motion_bin_length_s is set — they drive weighted start-time sampling even without the curriculum; the adaptive EMA layer activates only when adaptive_sampling=True, so record_bin_outcomes / blend_failure_ema / _update_adaptive_bin_weights keep their internal adaptive guards.

The term's start-time samplers (_sample_env_motion_start_time / _sample_as_concat_motion_bins) read the distribution via :attr:weights. Term env-assignment state (_assigned_env_motion_selection, _motion_buffer_start_time_s, envs_idx_to_assigned_ids) is read through the back-reference.

Classes

AdaptiveBinSampler

AdaptiveBinSampler(term: MotionTerm)

Owns the temporal bin weights and the adaptive start-bin curriculum.

Initialize temporal bin weights for weighted start-time sampling.

Source code in eden/managers/motion/bin_sampler.py

def __init__(self, term: MotionTerm) -> None:
    """Initialize temporal bin weights for weighted start-time sampling."""
    self._term = term
    device = term.device

    num_motion_bins = torch.clamp(
        torch.floor(
            (term.data_source.buffer_length / term.data_source.framerate).to(device) / term.motion_bin_length_s
        ),
        min=1,
    ).to(torch.long)
    self._motion_bin_weights = ConcatBatchTensor(
        batch_sizes=num_motion_bins.tolist(),
        data_shape=tuple(),
        device=device,
    )
    for motion_idx in range(len(term.data_source)):
        motion_bin_probability = torch.ones(num_motion_bins[motion_idx], device=device, dtype=torch.float)
        if motion_bin_probability.numel() > 0:
            motion_bin_probability /= motion_bin_probability.sum()
        self._motion_bin_weights[motion_idx] = motion_bin_probability

    # Per-death-bin failure-count EMA (see ``record_bin_outcomes``).
    # ``None`` (allocated below) unless ``adaptive_sampling`` — the bin
    # weights themselves are used for weighted start sampling regardless.
    self._bin_failure_ema: torch.Tensor | None = None
    self._current_bin_failed: torch.Tensor | None = None
    self._adaptive_kernel: torch.Tensor | None = None
    # Precomputed per-bin coordinates + ``max(sizes)`` for the vectorised
    # weight update; depend only on sizes/starts (immutable for this library).
    self._adaptive_motion_ids_flat: torch.Tensor | None = None
    self._adaptive_pos_in_motion_flat: torch.Tensor | None = None
    self._adaptive_max_bins: int = 0
    self._adaptive_window_idx: torch.Tensor | None = None
    self._group_bin_mask: torch.Tensor | None = None

    if term.adaptive_sampling:
        total_bins = self._motion_bin_weights.concatenated_tensor.shape[0]
        self._bin_failure_ema = torch.zeros(total_bins, device=device, dtype=torch.float)
        self._current_bin_failed = torch.zeros(total_bins, device=device, dtype=torch.float)
        # Cache the (motion_id, position-within-motion) coordinates and
        # ``max(sizes)`` used by ``_update_adaptive_bin_weights`` — they depend
        # only on ``batch_sizes`` / ``batch_starts`` (immutable after this
        # returns), so caching here removes per-call allocations and the
        # per-call ``sizes.max().item()`` GPU→CPU sync.
        sizes_t = self._motion_bin_weights.batch_sizes
        starts_t = self._motion_bin_weights.batch_starts
        n_motions = len(sizes_t)
        self._adaptive_motion_ids_flat = torch.repeat_interleave(torch.arange(n_motions, device=device), sizes_t)
        self._adaptive_pos_in_motion_flat = (
            torch.arange(total_bins, device=device) - starts_t[self._adaptive_motion_ids_flat]
        )
        self._adaptive_max_bins = int(sizes_t.max().item()) if n_motions > 0 else 0
        if term.adaptive_kernel_size > 1:
            self._adaptive_kernel = torch.tensor(
                [term.adaptive_lambda**i for i in range(term.adaptive_kernel_size)],
                device=device,
            )
            self._adaptive_kernel = self._adaptive_kernel / self._adaptive_kernel.sum()
            # Flat index of every motion's bins, right-padded with its own last bin to
            # ``max_bins + K - 1`` columns: the look-ahead window the smoothing correlates with.
            cols = torch.arange(self._adaptive_max_bins + term.adaptive_kernel_size - 1, device=device)
            self._adaptive_window_idx = starts_t[:, None] + torch.minimum(cols[None, :], sizes_t[:, None] - 1)

Attributes

group_bin_mask property

group_bin_mask: Tensor

(G, total_bins) membership of every start bin in each per-env motion group (env_motion_groups).

Built on first use from the term's per-rank (G, M) clip mask, and cached: the bin layout is fixed for this library, and the sampler is rebuilt whenever the library changes.

weights property

weights: ConcatBatchTensor

Per-motion temporal bin probabilities (read by the term's start-time samplers).

Methods:

record_bin_outcomes

record_bin_outcomes(envs_idx: Tensor, terminated_mask: Tensor, timestamps: Tensor)

Accumulate per-episode failure counts for resetting envs.

Failures are keyed by the bin the episode died in (floor((start_time + timestamps) / bin_length)) and accumulated as raw counts in _current_bin_failed for the terminated subset only. A failure rate keyed by the start bin instead starves the easy early frames a cold policy needs to bootstrap on multi-clip LAFAN1 and collapses tracking; the death-bin failure count is the curriculum signal that lets it learn.

terminated_mask / timestamps are aligned with envs_idx.

Source code in eden/managers/motion/bin_sampler.py

def record_bin_outcomes(self, envs_idx: torch.Tensor, terminated_mask: torch.Tensor, timestamps: torch.Tensor):
    """Accumulate per-episode failure counts for resetting envs.

    Failures are keyed by the bin the episode **died in**
    (``floor((start_time + timestamps) / bin_length)``) and accumulated as
    raw counts in ``_current_bin_failed`` for the terminated subset only.
    A failure *rate* keyed by the start bin instead starves the easy early
    frames a cold policy needs to bootstrap on multi-clip LAFAN1 and
    collapses tracking; the death-bin failure count is the curriculum
    signal that lets it learn.

    ``terminated_mask`` / ``timestamps`` are aligned with ``envs_idx``.
    """
    term = self._term
    if not term.adaptive_sampling or self._motion_bin_weights is None or len(envs_idx) == 0:
        return
    assigned_ids = term.envs_idx_to_assigned_ids(envs_idx)
    motion_ids = term._assigned_env_motion_selection[assigned_ids]
    death_time = term._motion_buffer_start_time_s[assigned_ids] + timestamps
    bin_indices = torch.clamp(torch.floor(death_time / term.motion_bin_length_s).long(), min=0)
    max_bins = self._motion_bin_weights.batch_sizes[motion_ids] - 1
    bin_indices = torch.min(bin_indices, max_bins)
    flat_indices = self._motion_bin_weights.batch_starts[motion_ids] + bin_indices
    # Every resetting env's bin, weighted 1 if it terminated and 0 if not: the same counts as scattering
    # ones over the terminated subset (whole numbers, so the order of the adds cannot matter), without
    # the bool-mask gather that picked the subset, a GPU->CPU sync.
    self._current_bin_failed.scatter_add_(0, flat_indices, terminated_mask.to(self._current_bin_failed.dtype))

blend_failure_ema

blend_failure_ema() -> None

Blend _current_bin_failed into _bin_failure_ema, then zero it.

Source code in eden/managers/motion/bin_sampler.py

def blend_failure_ema(self) -> None:
    """Blend ``_current_bin_failed`` into ``_bin_failure_ema``, then zero it."""
    term = self._term
    if not term.adaptive_sampling or self._bin_failure_ema is None:
        return
    self._bin_failure_ema = (
        term.adaptive_alpha * self._current_bin_failed + (1 - term.adaptive_alpha) * self._bin_failure_ema
    )
    self._current_bin_failed.zero_()

Per-motion-id outcome curriculum (hard-motion mining).

Split out of MotionTerm: owns the per-motion success/failure count buffers, the drain cadence, the weight-rewrite dynamics (snapshot / opt-in compound + uniform-mix regulariser), and the wandb summary/histogram reporting. One instance per term that sets adaptive_motion_weights=True, created in MotionTerm.__init__.

_motion_weights itself stays on the term — it is core sampling state read by _local_motion_weights / get_current_motion_weights. This curriculum rewrites it in place (through the back-reference) from the drained outcome rates; everything else it touches (counts, summaries) it owns.

Cadence: MotionTerm._update_motion_weight_curriculum calls update once per PPO iter; every motion_weight_update_every_n_iters-th call drains, optionally all-reduces across DDP, and rewrites the weights. The term piggybacks the per-motion metric drain on the same cadence.

Classes

MotionOutcomeCurriculum

MotionOutcomeCurriculum(term: MotionTerm)

Owns the per-motion-id outcome counts and the weight-rewrite pipeline.

Source code in eden/managers/motion/outcome_curriculum.py

def __init__(self, term: MotionTerm) -> None:
    self._term = term
    # Fail fast on bad curriculum config — caught here at term setup rather
    # than at the first drain (which can be many iters into training when
    # ``motion_weight_update_every_n_iters > 1``).
    if not (0.0 <= term.motion_weight_uniform_mix <= 1.0):
        raise ValueError(
            f"motion_weight_uniform_mix must be in [0, 1], got "
            f"{term.motion_weight_uniform_mix}. The regulariser blends "
            f"``w := (1 - α) * w + α / N`` after the curriculum update; "
            f"values outside [0, 1] produce non-distributions."
        )
    if term.motion_weight_update_mode == "snapshot" and term.motion_weight_failure_rate_max_over_mean <= 0.0:
        # Without a positive clamp ceiling, ``_apply_snapshot_curriculum``
        # produces an all-zero ``_motion_weights`` (clip_upper ≤ 0 → every
        # clipped rate is 0 → division by clamp_min(_NORMALIZE_EPS) yields
        # all zeros). The downstream ``torch.multinomial`` would crash or
        # return garbage. Reject the config at setup time.
        raise ValueError(
            f"motion_weight_failure_rate_max_over_mean must be > 0 in snapshot mode, "
            f"got {term.motion_weight_failure_rate_max_over_mean}. SONIC's default is "
            f"200.0 (lenient clamp); tighter values like 5.0 are valid and bite the "
            f"distribution earlier."
        )
    n = term._global_num_motions
    # Sized to the GLOBAL motion-id space so DDP all_reduce stays
    # shape-symmetric across ranks regardless of mp_split_method.
    self._motion_success_counts = torch.zeros(n, dtype=torch.long, device=term.device)
    self._motion_failure_counts = torch.zeros(n, dtype=torch.long, device=term.device)
    self._iters_since_motion_weight_update = 0
    # ``_curriculum_has_run`` gates the floor-clip metric: pre-first-drain
    # the user-supplied raw weights have never been clamped, so reporting a
    # "fraction at floor" relative to the configured floor is noise (e.g.
    # raw weights ``[1e-4, …]`` look 100% floor-clipped though the
    # curriculum hasn't touched them). Flip on once the first apply runs.
    self._curriculum_has_run = False
    # Cached summaries (mixed float/str; forwarded verbatim to wandb.log via
    # the rsl_rl mixin). Refreshed only on drain / init, so the hot reset
    # path is a dict lookup rather than 5+ GPU→CPU syncs.
    self._last_curriculum_summary: dict[str, float | str] = {}
    self._last_curriculum_per_motion: dict[str, np.ndarray] | None = None
    self._last_weights_per_motion: np.ndarray | None = None
    self._last_weights_summary: dict[str, float | str] = self._compute_motion_weights_summary()

Attributes

last_curriculum_summary property

last_curriculum_summary: dict[str, float | str]

Cached per-iter outcome scalars (MotionOutcome/*); empty until first drain.

Methods:

record_motion_outcomes

record_motion_outcomes(success_envs_idx: Tensor, failure_envs_idx: Tensor, success_timestamps: Tensor, failure_timestamps: Tensor, env_dt: float) -> None

Increment per-motion success / failure counts for envs that just reset.

Classification is driven by motion completion, not by which cohort an env arrived in (matches ProtoMotions's done_clip: a clip is "done" when motion_times + env_dt >= end_times, regardless of why the episode reset).

• Completed (start_time + episode_time + env_dt >= motion_length) → success. This includes terminated envs whose motion happened to play to its natural end.
• Terminated AND not completed → failure.
• Not terminated AND not completed (mid-motion timeout) → dropped.

Both *_envs_idx accept the same forms as envs_idx_to_assigned_ids: a bool mask of length num_envs (preferred — avoids GPU sync) or an int-index tensor of global env ids. *_timestamps are compact tensors aligned with the corresponding cohort.

Source code in eden/managers/motion/outcome_curriculum.py

def record_motion_outcomes(
    self,
    success_envs_idx: torch.Tensor,
    failure_envs_idx: torch.Tensor,
    success_timestamps: torch.Tensor,
    failure_timestamps: torch.Tensor,
    env_dt: float,
) -> None:
    """Increment per-motion success / failure counts for envs that just reset.

    Classification is driven by **motion completion**, not by which cohort
    an env arrived in (matches ProtoMotions's ``done_clip``: a clip is
    "done" when ``motion_times + env_dt >= end_times``, regardless of why
    the episode reset).

    - Completed (``start_time + episode_time + env_dt >= motion_length``)
      → success. This includes terminated envs whose motion happened to
      play to its natural end.
    - Terminated AND not completed → failure.
    - Not terminated AND not completed (mid-motion timeout) → dropped.

    Both ``*_envs_idx`` accept the same forms as ``envs_idx_to_assigned_ids``:
    a bool mask of length ``num_envs`` (preferred — avoids GPU sync) or an
    int-index tensor of global env ids. ``*_timestamps`` are compact tensors
    aligned with the corresponding cohort.
    """
    # Non-terminated cohort: only completed motions count as success; the
    # rest (mid-motion timeouts) are ambiguous and dropped on the floor.
    self._tally_outcomes(success_envs_idx, success_timestamps, env_dt, count_incomplete_as_failure=False)
    # Terminated cohort: completed motions still count as success (the
    # termination just coincided with natural end); the rest are failures.
    self._tally_outcomes(failure_envs_idx, failure_timestamps, env_dt, count_incomplete_as_failure=True)

update

update(all_reduce_fn=None) -> bool

Drain outcome counts, optionally reduce across DDP ranks, rewrite weights.

Called once per PPO iteration via MotionTerm._update_motion_weight_curriculum. The first motion_weight_update_every_n_iters - 1 calls just bump the internal iter counter and return False (counts keep accumulating); the K-th call drains, optionally all-reduces, and applies the curriculum, returning True (which the term uses to gate the metric drain onto the same cadence).

all_reduce_fn — if provided, applied in-place to both count tensors before they drive the curriculum (the manager supplies a dist.all_reduce lambda for DDP; tests pass None).

Source code in eden/managers/motion/outcome_curriculum.py

def update(self, all_reduce_fn=None) -> bool:
    """Drain outcome counts, optionally reduce across DDP ranks, rewrite weights.

    Called once per PPO iteration via ``MotionTerm._update_motion_weight_curriculum``.
    The first ``motion_weight_update_every_n_iters - 1`` calls just bump the
    internal iter counter and return ``False`` (counts keep accumulating);
    the K-th call drains, optionally all-reduces, and applies the curriculum,
    returning ``True`` (which the term uses to gate the metric drain onto the
    same cadence).

    ``all_reduce_fn`` — if provided, applied in-place to both count tensors
    before they drive the curriculum (the manager supplies a
    ``dist.all_reduce`` lambda for DDP; tests pass ``None``).
    """
    self._iters_since_motion_weight_update += 1
    if self._iters_since_motion_weight_update < self._term.motion_weight_update_every_n_iters:
        return False
    self._iters_since_motion_weight_update = 0

    success = self._motion_success_counts.clone()
    failure = self._motion_failure_counts.clone()
    self._motion_success_counts.zero_()
    self._motion_failure_counts.zero_()

    if all_reduce_fn is not None:
        all_reduce_fn(success)
        all_reduce_fn(failure)

    self._record_curriculum_summary(success, failure)
    self._apply_motion_outcome_curriculum(success, failure)
    self._curriculum_has_run = True
    # Refresh the weights-summary cache now that ``_motion_weights`` has been
    # rewritten. Done here (not on every reset) so the hot reset path doesn't
    # pay 5 GPU→CPU syncs every time an env terminates.
    self._last_weights_summary = self._compute_motion_weights_summary()
    return True

summarize_motion_weights

summarize_motion_weights() -> dict[str, float | str]

Read the cached scalar summary of the motion-weights distribution.

Returns the dict populated by _compute_motion_weights_summary, refreshed only when _motion_weights actually changes (init, plus each update call that drains). The hot reset path reads this cache rather than recomputing from the GPU tensor (which would cost 5 GPU→CPU syncs).

Source code in eden/managers/motion/outcome_curriculum.py

def summarize_motion_weights(self) -> dict[str, float | str]:
    """Read the cached scalar summary of the motion-weights distribution.

    Returns the dict populated by ``_compute_motion_weights_summary``,
    refreshed only when ``_motion_weights`` actually changes (init, plus each
    ``update`` call that drains). The hot reset path reads this cache rather than
    recomputing from the GPU tensor (which would cost 5 GPU→CPU syncs).
    """
    return self._last_weights_summary

collect_per_motion_distributions

collect_per_motion_distributions(wandb_module) -> dict[str, Any]

Build per-motion histograms for the wandb-direct emit path.

Returns wandb.Histogram objects keyed under MotionPerMotion/* when log_per_motion_metrics=True; an empty dict otherwise. Built from the cached numpy arrays _compute_motion_weights_summary / _record_curriculum_summary populate — one snapshot per drain, not per env step. Histograms (not per-id scalars) keep the fan-out constant-cost at AMASS / SEED scale.

Source code in eden/managers/motion/outcome_curriculum.py

def collect_per_motion_distributions(self, wandb_module) -> dict[str, Any]:
    """Build per-motion histograms for the wandb-direct emit path.

    Returns ``wandb.Histogram`` objects keyed under ``MotionPerMotion/*`` when
    ``log_per_motion_metrics=True``; an empty dict otherwise. Built from the
    cached numpy arrays ``_compute_motion_weights_summary`` /
    ``_record_curriculum_summary`` populate — one snapshot per drain, not per
    env step. Histograms (not per-id scalars) keep the fan-out constant-cost
    at AMASS / SEED scale.
    """
    if not self._term.log_per_motion_metrics:
        return {}
    out: dict[str, Any] = {}
    if self._last_weights_per_motion is not None:
        out["MotionPerMotion/Weights"] = wandb_module.Histogram(self._last_weights_per_motion)
    if self._last_curriculum_per_motion is not None:
        for field, arr in self._last_curriculum_per_motion.items():
            out[f"MotionPerMotion/Outcome/{field}"] = wandb_module.Histogram(arr)
    return out

Per-motion tracking-error metrics (MPJPE / G-MPJPE / root error).

Split out of MotionTerm: a self-contained collaborator owning the metric buffers + link-index cache and the record → scatter → drain → summary pipeline. One instance per opted-in term, created in MotionTerm.set_envs_idx_assignments when compute_per_motion_metrics is set; it reads the term's env-assignment / DDP id state through a back-reference (the term stays the owner of that state).

Pipeline (and why each stage runs when it does):

• record_step — every compute() step: per-env error vs the reference, added into the per-env running sums. No CPU↔GPU sync.
• scatter_at_reset — at episode end: mean each finished env's sums over its episode and scatter into the per-motion accumulators, sized to the GLOBAL motion-id space so DDP all_reduce stays shape-symmetric (mirrors the curriculum counts).
• drain — once per curriculum drain: reduce across DDP and cache the aggregate summary surfaced via MotionManager._collect_log_extras. Shares the curriculum's all_reduce_fn so it pays one extra collective.

Classes

PerMotionMetrics

PerMotionMetrics(term: MotionTerm)

Owns the per-motion tracking-error buffers and their accumulation pipeline.

Source code in eden/managers/motion/per_motion_metrics.py

def __init__(self, term: MotionTerm) -> None:
    self._term = term
    n_envs = term.num_assigned_envs
    n_motions = term._global_num_motions
    device = term.device
    # Lazily-resolved link-index caches; populated on the first ``record_step``.
    self._metric_robot_link_idx: torch.Tensor | None = None
    self._metric_motion_link_idx: torch.Tensor | None = None
    self._metric_robot_anchor_idx: torch.Tensor | None = None
    self._metric_motion_anchor_idx: torch.Tensor | None = None
    # Aggregate summary cache; refreshed only on drain.
    self._last_metrics_summary: dict[str, float] = {}
    # Per-env buffers are sized to the rank's local assigned slice; per-motion
    # buffers to the GLOBAL motion-id space (DDP all_reduce shape-symmetry).
    self._per_env_mpjpe_sum = torch.zeros(n_envs, dtype=torch.float, device=device)
    self._per_env_root_pos_err_sum = torch.zeros(n_envs, dtype=torch.float, device=device)
    self._per_env_root_quat_err_sum = torch.zeros(n_envs, dtype=torch.float, device=device)
    self._motion_mpjpe_sum = torch.zeros(n_motions, dtype=torch.float, device=device)
    self._motion_root_pos_err_sum = torch.zeros(n_motions, dtype=torch.float, device=device)
    self._motion_root_quat_err_sum = torch.zeros(n_motions, dtype=torch.float, device=device)
    self._motion_metric_episode_count = torch.zeros(n_motions, dtype=torch.long, device=device)
    # G-MPJPE buffers exist iff an anchor is configured; a single ``is None``
    # check on ``_per_env_g_mpjpe_sum`` gates the path.
    self._per_env_g_mpjpe_sum: torch.Tensor | None = None
    self._motion_g_mpjpe_sum: torch.Tensor | None = None
    if term.metric_anchor_link_name is not None:
        self._per_env_g_mpjpe_sum = torch.zeros(n_envs, dtype=torch.float, device=device)
        self._motion_g_mpjpe_sum = torch.zeros(n_motions, dtype=torch.float, device=device)

Methods:

record_step

record_step(entity: Entity, ref: MotionReferenceData, env_name: str) -> None

Compute one step of per-env tracking error for entity and add it into the running sums.

Per-entity no-op when the reference lacks world-frame link/root fields (e.g. dof-only motion data) or the term tracks no links for this entity. A get_links_pos (two with an anchor) + one get_pos/get_quat and per-env reductions, all on the GPU with no sync.

Source code in eden/managers/motion/per_motion_metrics.py

def record_step(self, entity: Entity, ref: MotionReferenceData, env_name: str) -> None:
    """Compute one step of per-env tracking error for ``entity`` and add it
    into the running sums.

    Per-entity no-op when the reference lacks world-frame link/root fields
    (e.g. dof-only motion data) or the term tracks no links for this entity.
    A ``get_links_pos`` (two with an anchor) + one ``get_pos``/``get_quat``
    and per-env reductions, all on the GPU with no sync.
    """
    term = self._term
    if ref.links_pos_w is None or ref.base_pos_w is None or ref.base_quat_w is None:
        return
    motion_key = term.entities_map.get(env_name, env_name)
    links_name = term.target_links_name.get(motion_key, [])
    if not links_name:
        return

    # Lazy-cache link indices on first call. Motion-side uses a contiguous
    # arange because ``data_source.select_links`` already reordered the ref
    # tensor to match ``target_links_name``.
    if self._metric_robot_link_idx is None:
        _, robot_idx = entity.find_named_links_idx_local(links_name, preserve_order=True)
        self._metric_robot_link_idx = torch.as_tensor(robot_idx, device=term.device, dtype=torch.long)
        self._metric_motion_link_idx = torch.arange(len(links_name), device=term.device, dtype=torch.long)
        if term.metric_anchor_link_name is not None:
            anchor_name = term.metric_anchor_link_name
            if anchor_name not in links_name:
                raise ValueError(
                    f"metric_anchor_link_name={anchor_name!r} must be present in "
                    f"target_links_name for entity {env_name!r}; got {links_name}"
                )
            _, robot_anchor = entity.find_named_links_idx_local([anchor_name], preserve_order=True)
            self._metric_robot_anchor_idx = torch.as_tensor(robot_anchor, device=term.device, dtype=torch.long)
            self._metric_motion_anchor_idx = torch.tensor(
                [links_name.index(anchor_name)], device=term.device, dtype=torch.long
            )

    # World-frame link MPJPE.
    robot_links = entity.get_links_pos(ls_idx_local=self._metric_robot_link_idx)
    motion_links = ref.links_pos_w[:, 0].index_select(1, self._metric_motion_link_idx)
    mpjpe = torch.linalg.vector_norm(robot_links - motion_links, dim=-1).mean(dim=-1)

    # G-MPJPE: link MPJPE relative to the anchor link's position — removes
    # the anchor's translation error (heading error still counts). Only
    # computed when an anchor is configured.
    g_mpjpe = None
    if self._metric_robot_anchor_idx is not None:
        robot_anchor_pos = entity.get_links_pos(ls_idx_local=self._metric_robot_anchor_idx)
        motion_anchor_pos = ref.links_pos_w[:, 0].index_select(1, self._metric_motion_anchor_idx)
        g_mpjpe = torch.linalg.vector_norm(
            (robot_links - robot_anchor_pos) - (motion_links - motion_anchor_pos),
            dim=-1,
        ).mean(dim=-1)

    # Root error: L2 distance + quat geodesic angle 2*acos(|q1·q2|).
    # ``clamp`` guards the |dot|>1 overshoot from float roundoff.
    root_pos_err = torch.linalg.vector_norm(entity.get_pos() - ref.base_pos_w[:, 0], dim=-1)
    q_dot = (entity.get_quat() * ref.base_quat_w[:, 0]).sum(dim=-1).abs().clamp(max=1.0)
    root_quat_err = 2.0 * torch.acos(q_dot)

    sl = term.assigned_env_slice
    self.accumulate(
        mpjpe=mpjpe[sl],
        root_pos_err=root_pos_err[sl],
        root_quat_err=root_quat_err[sl],
        g_mpjpe=g_mpjpe[sl] if g_mpjpe is not None else None,
    )

accumulate

accumulate(mpjpe: Tensor, root_pos_err: Tensor, root_quat_err: Tensor, g_mpjpe: Tensor | None = None) -> None

Add one step of per-env error into the running sums.

Inputs are [num_assigned_envs] float tensors on the term's device. g_mpjpe must be passed iff metric_anchor_link_name is set on the term — a mismatch is a config bug and raises.

Source code in eden/managers/motion/per_motion_metrics.py

def accumulate(
    self,
    mpjpe: torch.Tensor,
    root_pos_err: torch.Tensor,
    root_quat_err: torch.Tensor,
    g_mpjpe: torch.Tensor | None = None,
) -> None:
    """Add one step of per-env error into the running sums.

    Inputs are ``[num_assigned_envs]`` float tensors on the term's device.
    ``g_mpjpe`` must be passed iff ``metric_anchor_link_name`` is set on the
    term — a mismatch is a config bug and raises.
    """
    self._per_env_mpjpe_sum += mpjpe
    self._per_env_root_pos_err_sum += root_pos_err
    self._per_env_root_quat_err_sum += root_quat_err
    if (g_mpjpe is None) != (self._per_env_g_mpjpe_sum is None):
        raise ValueError(
            "g_mpjpe / _per_env_g_mpjpe_sum mismatch: pass g_mpjpe iff metric_anchor_link_name is set on the term."
        )
    if g_mpjpe is not None:
        self._per_env_g_mpjpe_sum += g_mpjpe

scatter_at_reset

scatter_at_reset(envs_idx: Tensor, episode_lengths: Tensor) -> None

Drain per-env running sums into the per-motion accumulators at episode end.

episode_lengths is the full-size per-env episode_length_buf ([num_envs]); we slice it to this term's assigned_env_slice FIRST so all downstream indexing operates in the LOCAL assigned-env space — envs_idx_to_assigned_ids returns local-space indices, so indexing the global episode_lengths with them silently reads the wrong env's length when assigned_env_slice.start > 0.

Supports both int-index and bool-mask envs_idx (the contract of record_motion_outcomes). Envs at their bootstrap reset (length 0) are filtered out — they have no real outcome. Maps motion ids through _local_to_global so per-rank contributions land in shape-compatible tensors for the downstream DDP all_reduce.

Source code in eden/managers/motion/per_motion_metrics.py

def scatter_at_reset(self, envs_idx: torch.Tensor, episode_lengths: torch.Tensor) -> None:
    """Drain per-env running sums into the per-motion accumulators at episode end.

    ``episode_lengths`` is the full-size per-env ``episode_length_buf``
    (``[num_envs]``); we slice it to this term's ``assigned_env_slice`` FIRST
    so all downstream indexing operates in the LOCAL assigned-env space —
    ``envs_idx_to_assigned_ids`` returns local-space indices, so indexing the
    global ``episode_lengths`` with them silently reads the wrong env's
    length when ``assigned_env_slice.start > 0``.

    Supports both int-index and bool-mask ``envs_idx`` (the contract of
    ``record_motion_outcomes``). Envs at their bootstrap reset (length 0) are
    filtered out — they have no real outcome. Maps motion ids through
    ``_local_to_global`` so per-rank contributions land in shape-compatible
    tensors for the downstream DDP ``all_reduce``.
    """
    term = self._term
    episode_lengths_local = episode_lengths[term.assigned_env_slice.start : term.assigned_env_slice.stop]

    assigned = term.envs_idx_to_assigned_ids(envs_idx)
    # For bool input ``assigned`` is a fixed-length mask (len always
    # num_assigned_envs), so the empty-tensor shortcut only applies to int input.
    if not assigned.dtype == torch.bool and len(assigned) == 0:
        return

    if assigned.dtype == torch.bool:
        assigned = env_indices(assigned)

    # Only envs that actually ran. ``episode_length_buf == 0`` is a bootstrap
    # reset — no real data to record. Selected with `torch.where` / 0-1 weights over every resetting
    # env rather than gathered: an `.any()` gate and a bool-mask gather are a GPU->CPU sync each, and
    # the counts and the zero contributions it scatters are exact.
    active_lengths = episode_lengths_local[assigned]
    ran_mask = active_lengths > 0
    lengths_f = active_lengths.clamp(min=1).to(self._per_env_mpjpe_sum.dtype)

    def episode_mean(per_env_sum: torch.Tensor) -> torch.Tensor:
        return torch.where(ran_mask, per_env_sum[assigned] / lengths_f, 0.0)

    def drain_ran(per_env_sum: torch.Tensor) -> None:
        per_env_sum[assigned] = torch.where(ran_mask, 0.0, per_env_sum[assigned])

    # Episode means (per active assigned env).
    mpjpe_mean = episode_mean(self._per_env_mpjpe_sum)
    root_pos_mean = episode_mean(self._per_env_root_pos_err_sum)
    root_quat_mean = episode_mean(self._per_env_root_quat_err_sum)

    motion_ids_local = term._assigned_env_motion_selection[assigned]
    motion_ids_global = term._local_to_global[motion_ids_local]

    self._motion_mpjpe_sum.scatter_add_(0, motion_ids_global, mpjpe_mean)
    self._motion_root_pos_err_sum.scatter_add_(0, motion_ids_global, root_pos_mean)
    self._motion_root_quat_err_sum.scatter_add_(0, motion_ids_global, root_quat_mean)
    self._motion_metric_episode_count.scatter_add_(
        0, motion_ids_global, ran_mask.to(self._motion_metric_episode_count.dtype)
    )

    # Zero per-env running sums for envs we just drained.
    drain_ran(self._per_env_mpjpe_sum)
    drain_ran(self._per_env_root_pos_err_sum)
    drain_ran(self._per_env_root_quat_err_sum)

    # G-MPJPE scatter is conditional on the anchor being configured.
    if self._per_env_g_mpjpe_sum is not None:
        self._motion_g_mpjpe_sum.scatter_add_(0, motion_ids_global, episode_mean(self._per_env_g_mpjpe_sum))
        drain_ran(self._per_env_g_mpjpe_sum)

drain

drain(all_reduce_fn=None) -> None

Drain per-motion accumulators, reduce across DDP, refresh summary cache.

Called from MotionTerm._update_motion_weight_curriculum right after the success/failure drain so both share one DDP cadence. After this returns, summary() holds the aggregate scalars surfaced via MotionManager._collect_log_extras.

Source code in eden/managers/motion/per_motion_metrics.py

def drain(self, all_reduce_fn=None) -> None:
    """Drain per-motion accumulators, reduce across DDP, refresh summary cache.

    Called from ``MotionTerm._update_motion_weight_curriculum`` right after
    the success/failure drain so both share one DDP cadence. After this
    returns, ``summary()`` holds the aggregate scalars surfaced via
    ``MotionManager._collect_log_extras``.
    """
    sums = {
        "mpjpe": self._motion_mpjpe_sum.clone(),
        "root_pos_err": self._motion_root_pos_err_sum.clone(),
        "root_quat_err": self._motion_root_quat_err_sum.clone(),
    }
    if self._motion_g_mpjpe_sum is not None:
        sums["g_mpjpe"] = self._motion_g_mpjpe_sum.clone()
        self._motion_g_mpjpe_sum.zero_()
    counts = self._motion_metric_episode_count.clone()
    self._motion_mpjpe_sum.zero_()
    self._motion_root_pos_err_sum.zero_()
    self._motion_root_quat_err_sum.zero_()
    self._motion_metric_episode_count.zero_()

    if all_reduce_fn is not None:
        for t in sums.values():
            all_reduce_fn(t)
        all_reduce_fn(counts)

    self._record_summary(sums, counts)

summary

summary() -> dict[str, float]

Aggregate scalar summary cached at the last drain (empty before then).

Source code in eden/managers/motion/per_motion_metrics.py

def summary(self) -> dict[str, float]:
    """Aggregate scalar summary cached at the last ``drain`` (empty before then)."""
    return self._last_metrics_summary

per_motion_tensors

per_motion_tensors() -> dict[str, Tensor]

Per-motion-id tensors (CPU dict views) suitable for CSV dump.

Returns the raw per-motion sum tensors and the episode count; callers compute the per-motion mean as sum / count.clamp_min(1) once they hold the tensors locally. Tensors are copied to CPU so the caller can iterate without holding GPU memory.

Source code in eden/managers/motion/per_motion_metrics.py

def per_motion_tensors(self) -> dict[str, torch.Tensor]:
    """Per-motion-id tensors (CPU dict views) suitable for CSV dump.

    Returns the raw per-motion sum tensors *and* the episode count; callers
    compute the per-motion mean as ``sum / count.clamp_min(1)`` once they
    hold the tensors locally. Tensors are copied to CPU so the caller can
    iterate without holding GPU memory.
    """
    out = {
        "mpjpe_sum": self._motion_mpjpe_sum.detach().cpu().clone(),
        "root_pos_err_sum": self._motion_root_pos_err_sum.detach().cpu().clone(),
        "root_quat_err_sum": self._motion_root_quat_err_sum.detach().cpu().clone(),
        "episode_count": self._motion_metric_episode_count.detach().cpu().clone(),
    }
    if self._motion_g_mpjpe_sum is not None:
        out["g_mpjpe_sum"] = self._motion_g_mpjpe_sum.detach().cpu().clone()
    return out

Functions:

Motion-end resync (opt-in): keep an episode running past a clip's end.

When MotionManagerOptions.resync_on_motion_end is set, an env whose reference reaches the end of its clip is soft-reset mid-episode — a fresh start phase is resampled, its playback clock restarted, and the robot teleported onto the new reference start — WITHOUT ending the episode. This mirrors the clip-wrap behavior of MotionCommand (the beyondmimic example); the default (off) freezes the reference at the clip's last frame.

It is folded into the single per-step fill_data(all): :meth:pre_fill detects wrapped envs from the just-advanced clock (before the fill) and resamples them, so the one fill_data computes their new-phase reference (no second fill); :meth:post_fill then teleports them onto that freshly-filled reference. One instance per :class:~eden.managers.motion_manager.MotionManager, created in its __init__; it reads the manager's state through a back-reference.

Hot-path cost contract: no GPU->CPU sync. Whether any env wrapped is never read on the host: every step resamples and teleports through the wrap mask, every env's row with the mask picking the wrapped ones (torch.where / masked_fill_, and Genesis' masked forward kinematics), so a step where nothing wrapped writes nothing — it only runs the same few kernels over all envs. That fixed cost does not outweigh the sync it replaces: BeyondMimic at 4096 envs (L40S) stepped in 24.12 ms against 24.16 ms with the nonzero() gate. (MotionCommand keeps its .any() gate: its reset path adds pose noise and full-batch draws, ~1 ms.)

Classes

MotionEndResync

MotionEndResync(manager: MotionManager)

Owns the opt-in motion-end resync operation for a MotionManager.

Source code in eden/managers/motion/resync.py

def __init__(self, manager: MotionManager) -> None:
    self._mm = manager
    # RSI RESET-event terms, resolved lazily (the event manager is built after the
    # motion manager) but forced at build time via :meth:`validate` when enabled.
    self._rsi_terms: list | None = None
    # Env lookup hoisted out of the per-step post_fill.
    self._debug = os.environ.get("EDEN_RESYNC_DEBUG") == "1"

Attributes

enabled property

enabled: bool

Methods:

validate

validate() -> None

Build-time check (a post-build hook, once the event manager exists): resolve the RSI teleport terms now so a missing one raises when resync_on_motion_end is first honored — at env construction — instead of at the first mid-episode wrap. The resolution is cached, so the first wrap pays nothing.

Source code in eden/managers/motion/resync.py

def validate(self) -> None:
    """Build-time check (a post-build hook, once the event manager exists): resolve
    the RSI teleport terms now so a missing one raises when ``resync_on_motion_end``
    is first honored — at env construction — instead of at the first mid-episode
    wrap. The resolution is cached, so the first wrap pays nothing.
    """
    self._reference_init_terms()

pre_fill

pre_fill() -> Tensor | None

Before fill_data(all): detect wrapped envs and resample them so the single fill computes their new-phase reference. Returns the (num_envs,) wrap mask (to hand back to :meth:post_fill), or None when disabled. The mask is not read on the host, so it may select no env: everything below writes through it, which then writes nothing.

Source code in eden/managers/motion/resync.py

def pre_fill(self) -> torch.Tensor | None:
    """Before ``fill_data(all)``: detect wrapped envs and resample them so the
    single fill computes their new-phase reference. Returns the ``(num_envs,)``
    wrap mask (to hand back to :meth:`post_fill`), or ``None`` when disabled.
    The mask is not read on the host, so it may select no env: everything below
    writes through it, which then writes nothing.
    """
    if not self.enabled:
        return None
    wrap_mask = self._wrapped_mask()
    self._resample(wrap_mask)
    return wrap_mask

post_fill

post_fill(wrap_mask: Tensor) -> None

After fill_data(all): teleport the wrapped envs onto the freshly filled reference. With EDEN_RESYNC_DEBUG=1, first assert the pre-fill wrap detection matched fill_data's validity (no env still past-end).

Source code in eden/managers/motion/resync.py

def post_fill(self, wrap_mask: torch.Tensor) -> None:
    """After ``fill_data(all)``: teleport the wrapped envs onto the freshly
    filled reference. With ``EDEN_RESYNC_DEBUG=1``, first assert the pre-fill wrap
    detection matched ``fill_data``'s validity (no env still past-end).
    """
    if self._debug:
        for _e, _ref in self._mm._data.items():
            assert not (~_ref.validity[:, 0]).any(), "resync: env still past clip end after resample+fill"
    self._teleport(wrap_mask)

Options

MotionManagerOptions, MotionTermOptions and MotionGroupOptions are on the manager options page.

Data containers and lookup helpers

The motion data containers live in eden/utils/motion_reference/, which has no __init__.py, so the API generator cannot render them here. Read their docstrings in the source:

Module	Contents
motion_data.py	MotionSequence (one clip, [num_frames, ...] arrays), BatchMotionSequence, MotionReferenceData, MotionReferenceState
lookup.py	index resolution for tracking terms: get_motion_link_idx, get_motion_dofs_idx, current_motion_frame, assert_motion_field_present, get_anchor_state
utils.py	robot-vs-reference differences (get_link_position_distance, get_dofs_position_difference, ...) and interpolation (interpolate_motion, slerp_rows)
transforms.py	anchor-relative link transforms
augmentation_mixin.py, filter_mixin.py	SymmetricAugmentationMixin, FilterMixin
bipedal_ik.py	start-transition IK used in preprocessing

See also

• Motion tracking guide
• Datasets
• Tracking observations, rewards and terminations