mjlab.managers

Environment managers.

Classes

ManagerBase ManagerTermBase ManagerTermBaseCfg SceneEntityCfg ActionManager ActionTerm ActionTermCfg ObservationManager ObservationGroupCfg ObservationTermCfg RewardManager

RewardTermCfg TerminationManager TerminationTermCfg CommandManager NullCommandManager CommandTerm CommandTermCfg CurriculumManager NullCurriculumManager CurriculumTermCfg

EventManager EventMode EventTermCfg MetricsManager NullMetricsManager MetricsTermCfg RecorderManager NullRecorderManager RecorderTerm RecorderTermCfg

Base

class mjlab.managers.ManagerBase

Bases: ABC

Base class for all managers.

__init__(env: ManagerBasedRlEnv)

reset(env_ids: Tensor) → dict[str, Any]

Resets the manager and returns logging info for the current step.

class mjlab.managers.ManagerTermBase

Bases: object

__init__(env: ManagerBasedRlEnv)

reset(env_ids: Tensor | slice | None) → Any

Resets the manager term.

class mjlab.managers.ManagerTermBaseCfg

Base configuration for manager terms.

This is the base config for terms in observation, reward, termination, curriculum, and event managers. It provides a common interface for specifying a callable and its parameters.

The func field accepts either a function or a class:

Function-based terms are simpler and suitable for stateless computations:

RewardTermCfg(func=mdp.joint_torques_l2, weight=-0.01)

Class-based terms are instantiated with (cfg, env) and useful when you need to:

• Cache computed values at initialization (e.g., resolve regex patterns to indices)

• Maintain state across calls

• Perform expensive setup once rather than every call

class posture:
  def __init__(self, cfg: RewardTermCfg, env: ManagerBasedRlEnv):
    # Resolve std dict to tensor once at init
    self.std = resolve_std_to_tensor(cfg.params["std"], env)

  def __call__(self, env, **kwargs) -> torch.Tensor:
    # Use cached self.std
    return compute_posture_reward(env, self.std)

RewardTermCfg(func=posture, params={"std": {".*knee.*": 0.3}}, weight=1.0)

Class-based terms can optionally implement reset(env_ids) for per-episode state.

func: Any

The callable that computes this term’s value. Can be a function or a class. Classes are auto-instantiated with (cfg=term_cfg, env=env).

params: dict[str, Any]

Additional keyword arguments passed to func when called.

class mjlab.managers.SceneEntityCfg

Configuration for a scene entity that is used by the manager’s term.

This configuration allows flexible specification of entity components either by name or by ID. During resolution, it ensures consistency between names and IDs, and can optimize to slice(None) when all components are selected.

name: str

The name of the entity in the scene.

joint_names: str | tuple[str, ...] | None = None

Names of joints to include. Can be a single string or tuple.

joint_ids: list[int] | slice

IDs of joints to include. Can be a list or slice.

body_names: str | tuple[str, ...] | None = None

Names of bodies to include. Can be a single string or tuple.

body_ids: list[int] | slice

IDs of bodies to include. Can be a list or slice.

geom_names: str | tuple[str, ...] | None = None

Names of geometries to include. Can be a single string or tuple.

geom_ids: list[int] | slice

IDs of geometries to include. Can be a list or slice.

site_names: str | tuple[str, ...] | None = None

Names of sites to include. Can be a single string or tuple.

site_ids: list[int] | slice

IDs of sites to include. Can be a list or slice.

actuator_names: str | list[str] | None = None

Names of actuators to include. Can be a single string or list.

actuator_ids: list[int] | slice

IDs of actuators to include. Can be a list or slice.

tendon_names: str | tuple[str, ...] | None = None

Names of tendons to include. Can be a single string or tuple.

tendon_ids: list[int] | slice

IDs of tendons to include. Can be a list or slice.

camera_names: str | tuple[str, ...] | None = None

Names of cameras to include. Can be a single string or tuple.

camera_ids: list[int] | slice

IDs of cameras to include. Can be a list or slice.

light_names: str | tuple[str, ...] | None = None

Names of lights to include. Can be a single string or tuple.

light_ids: list[int] | slice

IDs of lights to include. Can be a list or slice.

material_names: str | tuple[str, ...] | None = None

Names of materials to include. Can be a single string or tuple.

material_ids: list[int] | slice

IDs of materials to include. Can be a list or slice.

pair_names: str | tuple[str, ...] | None = None

Names of contact pairs to include. Can be a single string or tuple.

pair_ids: list[int] | slice

IDs of contact pairs to include. Can be a list or slice.

preserve_order: bool = False

If True, maintains the order of components as specified.

resolve(scene: Scene) → None

Resolve names and IDs for all configured fields.

This method ensures consistency between names and IDs for each field type. It handles three cases: 1. Both names and IDs provided: Validates they match 2. Only names provided: Computes IDs (optimizes to slice(None) if all selected) 3. Only IDs provided: Computes names

Parameters:

scene – The scene containing the entity to resolve against.

Raises:

• ValueError – If provided names and IDs are inconsistent.

• KeyError – If the entity name is not found in the scene.

Action Manager

class mjlab.managers.ActionManager

Bases: ManagerBase

Manages action processing for the environment.

The action manager aggregates multiple action terms, each controlling a different entity or aspect of the simulation. It splits the policy’s action tensor and routes each slice to the appropriate action term.

__init__(cfg: dict[str, ActionTermCfg], env: ManagerBasedRlEnv)

property action: Tensor

Raw policy output from the current step, before per-term scale/offset. Shape: (num_envs, total_action_dim).

property prev_action: Tensor

Raw policy output from the previous step, before per-term scale/offset. Shape: (num_envs, total_action_dim).

property prev_prev_action: Tensor

Raw policy output from two steps ago, before per-term scale/offset. Shape: (num_envs, total_action_dim).

reset(env_ids: Tensor | slice | None = None) → dict[str, float]

Resets the manager and returns logging info for the current step.

process_action(action: Tensor) → None

Store the raw policy output and route slices to each action term.

Called once per policy step. The raw action tensor is saved into the history buffers (action, prev_action, prev_prev_action) before any per-term scale/offset is applied. Each term then receives its slice and independently applies its own affine transformation via ActionTerm.process_actions().

apply_action() → None

Write processed actions to entity actuator targets.

Called on every decimation substep (physics step), not just once per policy step. Each term writes its most recently processed targets to the simulation.

class mjlab.managers.ActionTerm

Bases: ManagerTermBase

Base class for action terms.

The action term is responsible for processing the raw actions sent to the environment and applying them to the entity managed by the term.

__init__(cfg: ActionTermCfg, env: ManagerBasedRlEnv)

class mjlab.managers.ActionTermCfg

Configuration for an action term.

Action terms process raw actions from the policy and apply them to entities in the scene (e.g., setting joint positions, velocities, or efforts).

entity_name: str

Name of the entity in the scene that this action term controls.

clip: dict[str, tuple] | None = None

Optional clipping bounds applied to processed actions (after scale and offset). Dict maps actuator name regex patterns to (min, max) tuples, resolved the same way as scale and offset.

abstractmethod build(env: ManagerBasedRlEnv) → ActionTerm

Build the action term from this config.

Observation Manager

class mjlab.managers.ObservationManager

Bases: ManagerBase

Manages observation computation for the environment.

The observation manager computes observations from multiple terms organized into groups. Each term can have noise, clipping, scaling, delay, and history applied. Groups can optionally concatenate their terms into a single tensor.

__init__(cfg: dict[str, ObservationGroupCfg], env)

reset(env_ids: Tensor | slice | None = None) → dict[str, float]

Resets the manager and returns logging info for the current step.

class mjlab.managers.ObservationGroupCfg

Configuration for an observation group.

An observation group bundles multiple observation terms together. Groups are typically used to separate observations for different purposes (e.g., “actor” for the actor, “critic” for the value function).

terms: dict[str, ObservationTermCfg]

Dictionary mapping term names to their configurations.

concatenate_terms: bool = True

Whether to concatenate all terms into a single tensor. If False, returns a dict mapping term names to their individual tensors.

concatenate_dim: int = -1

Dimension along which to concatenate terms. Default -1 (last dimension).

enable_corruption: bool = False

Whether to apply noise corruption to observations. Set to True during training for domain randomization, False during evaluation.

history_length: int | None = None

Group-level history length override. If set, applies to all terms in this group. If None, each term uses its own history_length setting.

flatten_history_dim: bool = True

Whether to flatten history into the observation dimension. If True, observations have shape (num_envs, obs_dim * history_length). If False, shape is (num_envs, history_length, obs_dim).

nan_policy: Literal['disabled', 'warn', 'sanitize', 'error'] = 'disabled'

NaN/Inf handling policy for observations in this group.

• ‘disabled’: No checks (default, fastest)

• ‘warn’: Log warning with term name and env IDs, then sanitize (debugging)

• ‘sanitize’: Silent sanitization to 0.0 like reward manager (safe for production)

• ‘error’: Raise ValueError on NaN/Inf (strict development mode)

nan_check_per_term: bool = True

If True, check each observation term individually to identify NaN source. If False, check only the final concatenated output (faster but less informative). Only applies when nan_policy != ‘disabled’.

class mjlab.managers.ObservationTermCfg

Configuration for an observation term.

Processing pipeline: compute → noise → clip → scale → delay → history. Delay models sensor latency. History provides temporal context. Both are optional and can be combined.

noise: NoiseCfg | NoiseModelCfg | None = None

Noise model to apply to the observation.

clip: tuple[float, float] | None = None

Range (min, max) to clip the observation values.

scale: tuple[float, ...] | float | Tensor | None = None

Scaling factor(s) to multiply the observation by.

delay_min_lag: int = 0

Minimum lag (in steps) for delayed observations. Lag sampled uniformly from [min_lag, max_lag]. Convert to ms: lag * (1000 / control_hz).

delay_max_lag: int = 0

Maximum lag (in steps) for delayed observations. Use min=max for constant delay.

delay_per_env: bool = True

If True, each environment samples its own lag. If False, all environments share the same lag at each step.

delay_hold_prob: float = 0.0

Probability of reusing the previous lag instead of resampling. Useful for temporally correlated latency patterns.

delay_update_period: int = 0

Resample lag every N steps (models multi-rate sensors). If 0, update every step.

delay_per_env_phase: bool = True

If True and update_period > 0, stagger update timing across envs to avoid synchronized resampling.

history_length: int = 0

Number of past observations to keep in history. 0 = no history.

flatten_history_dim: bool = True

Whether to flatten the history dimension into observation.

When True and concatenate_terms=True, uses term-major ordering: [A_t0, A_t1, …, A_tH-1, B_t0, B_t1, …, B_tH-1, …] See docs/source/observation.rst for details on ordering.

Reward Manager

class mjlab.managers.RewardManager

Bases: ManagerBase

Manages reward computation by aggregating weighted reward terms.

Reward Scaling Behavior:

By default, rewards are scaled by the environment step duration (dt). This normalizes cumulative episodic rewards across different simulation frequencies. The scaling can be disabled via the scale_by_dt parameter.

When scale_by_dt=True (default):

• reward_buf (returned by compute()) = raw_value * weight * dt

• _episode_sums (cumulative rewards) are scaled by dt

• Episode_Reward/* logged metrics are scaled by dt

When scale_by_dt=False:

• reward_buf = raw_value * weight (no dt scaling)

Regardless of the scaling setting:

• _step_reward (via get_active_iterable_terms()) always contains the unscaled reward rate (raw_value * weight)

__init__(cfg: dict[str, RewardTermCfg], env: ManagerBasedRlEnv, *, scale_by_dt: bool = True)

reset(env_ids: Tensor | slice | None = None) → dict[str, Tensor]

Resets the manager and returns logging info for the current step.

debug_vis(visualizer: DebugVisualizer) → None

Delegate debug visualization to class-based reward terms.

get_visualizable_terms() → list[tuple[str, Any]]

Return (name, func) pairs for class-based terms with debug_vis.

class mjlab.managers.RewardTermCfg

Configuration for a reward term.

func: Any

The callable that computes this reward term’s value.

weight: float

Weight multiplier for this reward term.

Termination Manager

class mjlab.managers.TerminationManager

Bases: ManagerBase

Manages termination conditions for the environment.

The termination manager aggregates multiple termination terms to compute episode done signals. Terms can be either truncations (time-based) or terminations (failure conditions).

__init__(cfg: dict[str, TerminationTermCfg], env: ManagerBasedRlEnv)

reset(env_ids: Tensor | slice | None = None) → dict[str, Tensor]

Resets the manager and returns logging info for the current step.

class mjlab.managers.TerminationTermCfg

Configuration for a termination term.

time_out: bool = False

Whether the term contributes towards episodic timeouts.

Command Manager

class mjlab.managers.CommandManager

Bases: ManagerBase

Manages command generation for the environment.

The command manager generates and updates goal commands for the agent (e.g., target velocity, target position). Commands are resampled at configurable intervals and can track metrics for logging.

__init__(cfg: dict[str, CommandTermCfg], env: ManagerBasedRlEnv)

create_gui(server: viser.ViserServer, get_env_idx: Callable[[], int], on_change: Callable[[], None] | None = None, request_action: Callable[[str, Any], None] | None = None) → None

Let each command term create its GUI controls.

on_viewer_pause(paused: bool) → None

Notify all command terms of viewer pause state change.

apply_gui_reset(env_ids: Tensor) → bool

Apply GUI-selected state from all terms. Returns True if any applied.

create_debug_vis_gui(server: viser.ViserServer, on_change: Callable[[], None] | None = None) → None

Add per-term debug visualization checkboxes.

reset(env_ids: Tensor | None) → dict[str, Tensor]

Resets the manager and returns logging info for the current step.

class mjlab.managers.NullCommandManager

Bases: object

Placeholder for absent command manager that safely no-ops all operations.

__init__()

class mjlab.managers.CommandTerm

Bases: ManagerTermBase

Base class for command terms.

__init__(cfg: CommandTermCfg, env: ManagerBasedRlEnv)

create_gui(name: str, server: viser.ViserServer, get_env_idx: Callable[[], int], on_change: Callable[[], None] | None = None, request_action: Callable[[str, Any], None] | None = None) → None

Create interactive GUI controls for this command term.

Override in subclasses to add task-specific controls (e.g., velocity sliders) to the Viser viewer. Called once during viewer setup.

The name argument is the term’s key in the command manager config (e.g., "twist").

on_viewer_pause(paused: bool) → None

Called when the viewer pause state changes.

apply_gui_reset(env_ids: Tensor) → bool

Apply GUI-selected state as an env reset override.

Returns True if this term wrote state to sim.

reset(env_ids: Tensor | slice | None) → dict[str, float]

Resets the manager term.

class mjlab.managers.CommandTermCfg

Configuration for a command generator term.

Command terms generate goal commands for the agent (e.g., target velocity, target position). Commands are automatically resampled at configurable intervals and can track metrics for logging.

resampling_time_range: tuple[float, float]

Time range in seconds for command resampling. When the timer expires, a new command is sampled and the timer is reset to a value uniformly drawn from [min, max]. Set both values equal for fixed-interval resampling.

debug_vis: bool = False

Whether to enable debug visualization for this command term. When True, the command term’s _debug_vis_impl method is called each frame to render visual aids (e.g., velocity arrows, target markers).

abstractmethod build(env: ManagerBasedRlEnv) → CommandTerm

Build the command term from this config.

Curriculum Manager

class mjlab.managers.CurriculumManager

Bases: ManagerBase

Manages curriculum learning for the environment.

The curriculum manager updates environment parameters during training based on agent performance. Each term can modify different aspects of the task difficulty (e.g., terrain complexity, command ranges).

__init__(cfg: dict[str, CurriculumTermCfg], env: ManagerBasedRlEnv)

reset(env_ids: Tensor | slice | None = None) → dict[str, float]

Resets the manager and returns logging info for the current step.

class mjlab.managers.NullCurriculumManager

Bases: object

Placeholder for absent curriculum manager that safely no-ops all operations.

__init__()

class mjlab.managers.CurriculumTermCfg

Configuration for a curriculum term.

Curriculum terms modify environment parameters during training to implement curriculum learning strategies (e.g., gradually increasing task difficulty).

Event Manager

class mjlab.managers.EventManager

Bases: ManagerBase

Manages event-based operations for the environment.

The event manager triggers operations at different simulation events: startup (once at initialization), reset (on episode reset), or interval (periodically during simulation). Common uses include domain randomization and state resets.

__init__(cfg: dict[str, EventTermCfg], env: ManagerBasedRlEnv)

get_term_cfg(term_name: str) → EventTermCfg

Get the configuration of a specific event term by name.

reset(env_ids: Tensor | None = None)

Resets the manager and returns logging info for the current step.

debug_vis(visualizer: DebugVisualizer) → None

Delegate debug visualization to class-based event terms.

mjlab.managers.EventMode

alias of Literal[‘startup’, ‘reset’, ‘interval’, ‘step’]

class mjlab.managers.EventTermCfg

Configuration for an event term.

Event terms trigger operations at specific simulation events. They’re commonly used for domain randomization, state resets, and periodic perturbations.

The four modes determine when the event fires:

• "startup": Once when the environment initializes. Use for parameters that should be randomized per-environment but stay constant within an episode (e.g., domain randomization).

• "reset": On every episode reset. Use for parameters that should vary between episodes (e.g., initial robot pose, domain randomization).

• "interval": Periodically during simulation, controlled by interval_range_s. Use for perturbations that should happen during episodes (e.g., pushing the robot, external disturbances).

• "step": Every environment step, unconditionally on all envs. Use for terms that manage per-step state such as force lifetimes (e.g., apply_body_impulse).

mode: EventMode

"startup" (once at init), "reset" (every episode), "interval" (periodically during simulation), or "step" (every environment step).

Type:

When the event triggers

interval_range_s: tuple[float, float] | None = None

Time range in seconds for interval mode. The next trigger time is uniformly sampled from [min, max]. Required when mode="interval".

is_global_time: bool = False

Whether all environments share the same timer. If True, all envs trigger simultaneously. If False (default), each env has an independent timer that resets on episode reset. Only applies to mode="interval".

min_step_count_between_reset: int = 0

Minimum environment steps between triggers. Prevents the event from firing too frequently when episodes reset rapidly. Only applies to mode="reset". Set to 0 (default) to trigger on every reset.

Metrics Manager

class mjlab.managers.MetricsManager

Bases: ManagerBase

Accumulates per-step metric values, reports episode averages.

Unlike rewards, metrics have no weight, no dt scaling, and no normalization by episode length. Episode values are true per-step averages (sum / step_count), so a metric in [0,1] stays in [0,1] in the logger.

__init__(cfg: dict[str, MetricsTermCfg], env: ManagerBasedRlEnv)

reset(env_ids: Tensor | slice | None = None) → dict[str, Tensor]

Resets the manager and returns logging info for the current step.

compute_substep() → None

Accumulate per-substep metric values inside the decimation loop.

No-op when no per_substep terms are configured.

class mjlab.managers.NullMetricsManager

Bases: object

Placeholder for absent metrics manager that safely no-ops all operations.

__init__()

class mjlab.managers.MetricsTermCfg

Configuration for a metrics term.

per_substep

If True, evaluate this term once per physics substep inside the

Type:

bool

decimation loop and report the per-step mean. Only the integrated state

(qpos, qvel, act) is current mid-loop; all derived quantities (xpos, xquat,

site_xpos, actuator_force, contacts, ...) are stale.

reduce

How to aggregate per-step values into an episode metric.

Type:

Literal[‘mean’, ‘last’]

``"mean"``

Type:

default

the value from the final step of the episode, which is useful for binary

success metrics that should not be averaged over timesteps.

Recorder Manager

class mjlab.managers.RecorderManager

Bases: ManagerBase

Orchestrates recorder terms during environment rollouts.

Holds a collection of RecorderTerm instances and calls their lifecycle methods at the appropriate points in the environment loop. The manager has no opinion on how data is stored; each term handles its own I/O entirely.

Register terms by adding them to the recorders dict on ManagerBasedRlEnvCfg. If the dict is empty, the environment substitutes a NullRecorderManager with zero overhead.

__init__(cfg: dict[str, RecorderTermCfg], env: ManagerBasedRlEnv)

record_pre_reset(env_ids: Tensor) → None

Forward to each term’s RecorderTerm.record_pre_reset().

record_post_reset(env_ids: Tensor) → None

Forward to each term’s RecorderTerm.record_post_reset().

record_post_step() → None

Forward to each term’s RecorderTerm.record_post_step().

close() → None

Forward to each term’s RecorderTerm.close().

class mjlab.managers.NullRecorderManager

Bases: object

No-op fallback used when no recorder terms are configured.

All methods are no-ops. This class is not a ManagerBase subclass so it carries zero overhead.

__init__()

class mjlab.managers.RecorderTerm

Bases: ManagerTermBase

Base class for recorder terms.

Override only the lifecycle methods you need. Each method is a no-op by default so subclasses are not required to implement all of them.

The environment is available as self._env, giving access to self._env.obs_buf, self._env.action_manager.action, and all other environment state.

Example:

class CsvRecorder(RecorderTerm):
  def __init__(self, cfg, env):
    super().__init__(cfg, env)
    self._file = open(cfg.params["path"], "w", newline="")
    self._writer = csv.writer(self._file)

  def record_pre_reset(self, env_ids):
    # Terminal transition: action is still intact here.
    # It will be zeroed by _reset_idx immediately after this returns.
    obs = self._env.obs_buf["actor"][env_ids].cpu().numpy()
    act = self._env.action_manager.action[env_ids].cpu().numpy()
    for o, a in zip(obs, act):
      self._writer.writerow(o.tolist() + a.tolist())

  def record_post_step(self):
    # Skip envs that just reset: their terminal pair was written in record_pre_reset
    # and their action is now zeroed.
    mask = ~self._env.reset_buf
    obs = self._env.obs_buf["actor"][mask].cpu().numpy()
    act = self._env.action_manager.action[mask].cpu().numpy()
    for o, a in zip(obs, act):
      self._writer.writerow(o.tolist() + a.tolist())

  def close(self):
    self._file.close()

__init__(cfg: RecorderTermCfg, env: ManagerBasedRlEnv)

record_pre_reset(env_ids: Tensor) → None

Called in env.step() before terminated environments are reset.

What is available:

• obs_buf contains the observation from the end of the previous step (the input the agent used to choose the terminal action). It does not contain the post-action terminal observation (the state reached after applying the action), which is never computed for resetting environments.

• action_manager.action contains the action applied during this step. This is the correct terminal action. It will be zeroed for these environments by _reset_idx immediately after this hook returns, so capture it here if you need it later.

• reward_buf contains the reward for this terminal step.

• reset_terminated and reset_time_outs reflect why each environment is resetting.

This is the right hook to record the terminal transition (obs_t, action_t, reward_t, done=True) for each resetting environment.

Parameters:

env_ids – Indices of environments that are about to be reset.

record_post_reset(env_ids: Tensor) → None

Called after a reset completes with fresh observations computed.

Fires at the end of env.reset() (covering all environments on the initial call) and within env.step() for each batch of environments that terminates, after state has been overwritten and new observations computed.

At this point obs_buf[env_ids] holds the initial observation of the new episode and action_manager.action[env_ids] is zero (no action has been taken in the new episode yet).

Use this hook to initialize per-episode state or record the first observation of a new episode.

Parameters:

env_ids – Indices of environments that were reset.

record_post_step() → None

Called at the end of every env.step() with fresh observations.

At this point obs_buf holds the new observation for every environment and action_manager.action holds the action that was applied during this step. Exception: for environments that reset during this step, action_manager.action has been zeroed by _reset_idx and obs_buf holds the initial observation of the new episode rather than the post-action terminal observation. Use record_pre_reset to capture the terminal (obs, action) pair for those environments. Resetting environments are identified by self._env.reset_buf.

close() → None

Called when the environment closes.

Release file handles, flush write buffers, or finalize output here.

class mjlab.managers.RecorderTermCfg

Configuration for a recorder term.

func must be a RecorderTerm subclass. Function-based terms are not supported because recorder terms are stateful (file handles, buffers, etc.).