reaf-sim API#

reaf.simulation.torch#

class ActionHistory#

Tracks commands and exposes them as features.

Flow:

process_commands(): stores raw action, passes command through.
produce_features(): returns current raw action and (optionally) past raw actions.

Parameters:

name – Unique identifier for this component.
key – The command key to consume and pass through (e.g. "actions"). Also used as the feature key for the current raw action.
feature_shape – Shape of the action tensor per environment (excluding the batch dimension).
nworld – Number of parallel environments.
device – Torch device for tensor allocations.
history_length – Number of past actions to buffer. 0 means only the current action feature is produced (no history).
reset_action – Action to fill at reset with shape (nworld, *feature_shape). If None, defaults to zeros.

property name: str#

property reset_action: torch.Tensor#: Action tensor of shape (nworld, *feature_shape) filled on reset.

consumed_commands_spec() → Mapping[str, gdm_robotics.interfaces.types.AnyArraySpec]#

produced_commands_keys() → set[str]#

process_commands(consumed_commands: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#: Store raw action in history, pass command through unchanged.

produced_features_spec() → Mapping[str, dm_env.specs.Array]#

required_features_keys() → set[str]#

produce_features(required_features: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#: Return current raw action and (optionally) past raw action history.

reset(env_mask: torch.Tensor = None) → None#: Reset current action and history buffer for the given environments.

class ActuatorDevice#

Actuate joints, tendons, or sites using existing MJCF actuators.

property name: str#: Returns the instance name of this device.

property num_targets: int#

property home_ctrl: torch.Tensor#

(num_ctrl,).

Type:: Home control values. Shape

initialize(simulator: Simulator) → None#: Resolve joint, actuator, and tendon indices.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the device state in the simulation.

substep_commands_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep commands spec for this device.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep measurements spec for this device.

set_substep_commands(simulator: Simulator, consumed_substep_commands: Mapping[str, torch.Tensor]) → None#: Sets the substep commands for this device.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#

Returns the measurements for the current state.

Should not advance any state: DeviceLayer can call this method any number of times and it should return valid measurements at any point in the step lifecycle.

Parameters:: simulator – The simulator.
Returns:: The measurements, batched over worlds. They may be views of simulator or device state, so callers must not modify them in place, and must copy them to keep them across simulator updates.

set_ctrl(env_mask: torch.Tensor, ctrl: torch.Tensor) → None#: Set actuator control for specified environments.

class ArticulatedDevice#

Manages the kinematic state of a robot subtree.

property name: str#: Returns the instance name of this device.

property home_qpos: torch.Tensor#

(nq,).

Type:: Home position for joints. Shape

property home_qvel: torch.Tensor#

(nv,).

Type:: Home velocity for joints. Shape

initialize(simulator: Simulator) → None#: Initializes the device with the compiled simulator.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the device state in the simulation.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep measurements spec for this device.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#

Returns the measurements for the current state.

Should not advance any state: DeviceLayer can call this method any number of times and it should return valid measurements at any point in the step lifecycle.

Parameters:: simulator – The simulator.
Returns:: The measurements, batched over worlds. They may be views of simulator or device state, so callers must not modify them in place, and must copy them to keep them across simulator updates.

set_qpos(env_mask: torch.Tensor, qpos: torch.Tensor) → None#: Set joint positions for specified environments.

set_qvel(env_mask: torch.Tensor, qvel: torch.Tensor) → None#: Set joint velocities for specified environments.

class BatchedMujocoCameraDevice#

Renders multiple cameras using Python mujoco.Renderer.

Creates its own render context with singleton pattern per environment.

property name: str#: Returns the instance name of this device.

initialize(simulator: Simulator) → None#: Initialize entity with compiled model - creates renderer.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the device state in the simulation.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep measurements spec for this device.

on_substep(simulator: Simulator) → None#: Renders on render ticks.

on_begin_stepping(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Renders the first frame after a reset.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#: Returns the latest rendered frames, without rendering.

render_now(simulator: Simulator, output_names: Sequence[str] | None = None, render_mode: RenderMode | None = None) → dict[str, torch.Tensor]#

Renders cameras immediately without affecting device measurements.

Side-channel call for Impl.CPP single-world environments that renders from the current simulator state and returns only the requested cameras and modes, without modifying cached substep measurements or the render_rate_hz schedule.

Parameters:

simulator – The simulator instance.
output_names – Optional subset of camera output_name`s to render. `None renders all cameras configured on this device.
render_mode – Optional subset of enabled RenderMode flags to render. None renders all modes configured on this device.

Returns:

A dictionary mapping {output_name}_{mode} to the freshly rendered tensor for each requested camera and mode.

Raises:

RuntimeError – If the device has not been initialized.
ValueError – If self._impl != types.Impl.CPP, data.nworld is greater than 1, output_names is a bare string, empty, or contains an unknown camera, or render_mode is empty or includes modes not enabled on this device.

substep_commands_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep commands spec for this device.

set_substep_commands(simulator: Simulator, substep_commands: Mapping[str, torch.Tensor]) → None#: Sets the substep commands for this device.

class CameraDevice#

Entity that renders the scene using MuJoCo rendering.

initialize(simulator: Simulator) → None#: Initializes the entity with compiled model and data.

render(model: TorchWrapper[mujoco_warp.Model] | None = None, data: TorchWrapper[mujoco_warp.Data] | None = None) → None#: Render all cameras (called once per scene.update).

property name: str#: Returns the instance name of this device.

property measurement_names: dict[str, str]#

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the processor.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Spec for measurements generated once per step.

on_substep(simulator: Simulator) → None#: Renders on render ticks.

on_begin_stepping(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Renders the first frame after a full reset.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#: Returns the latest rendered frames, without rendering.

class ClockDevice#

Provides the current simulation time at each step.

initialize(simulator: Simulator) → None#: Initialize entity with compiled model. ClockDevice has no state.

property name: str#: Returns the name of this entity.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: This entity has no state to reset.

substep_commands_spec() → Mapping[str, dm_env.specs.Array]#: This entity does not accept any commands.

set_substep_commands(simulator: Simulator, consumed_substep_commands: Mapping[str, torch.Tensor]) → None#: This entity does not process any commands.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#

Defines the structure of the time measurement.

Returns:: A dictionary specifying a single scalar float32 for the time.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#: Returns the current simulation time.

class ContactSensorDevice#

Contact sensors with pattern matching and configurable fields.

Uses MuJoCo naming: - objtype/objname: What to sense contact on - reftype/refname: What to sense contact against (optional) - fields: Which data to extract (found, force, etc.) - stacked together

configure_spec(spec: mujoco.MjSpec, *, objtype: Literal['geom', 'body', 'xbody'], objname: str | tuple[str, ...], reftype: Literal['geom', 'body', 'xbody'] | None = None, refname: str | None = None, reduce: Literal['none', 'mindist', 'maxforce', 'netforce'] = 'none', num: int = 1) → mujoco.MjSpec#

Add contact sensors to the spec via pattern matching.

Parameters:

spec – MjSpec to modify.
objtype – “geom”, “body”, or “xbody”.
objname – Name(s) to match on obj side (exact or regex pattern).
reftype – “geom”, “body”, or “xbody” (None = any).
refname – Name to match on ref side (None = any).
reduce – Reduction mode for multiple contacts.
num – How many contacts to report.

property name: str#: Returns the instance name of this device.

initialize(simulator: Simulator) → None#: Find sensors by prefix and build index tensor for fast gather.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the device state in the simulation.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Spec for contact sensor readings. Returns empty if dark.

on_substep(simulator: Simulator) → None#: Pushes the current contact reading into the substep history.

on_begin_stepping(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Fills the history of env_mask worlds with the reset reading.

get_measurements(simulator: Simulator, *, force: bool = False) → Mapping[str, torch.Tensor]#

Return contact sensor readings.

Parameters:

simulator – The simulator instance.
force – If True, return measurements even if dark (for resetter access).

class DefaultDiscountProvider#

Computes a constant discount given the termination state.

This provider returns a discount of 0.0 in case of termination and 1.0 otherwise (i.e. for truncation and not termination).

It is usually safe to use this discount provider for environments that return strictly positive rewards.

name() → str#

compute_discount(required_features: Mapping[str, torch.Tensor], termination_state: torch.Tensor) → torch.Tensor#: Computes the discount.

discount_spec() → dm_env.specs.Array#

required_features_keys() → set[str]#

class DelayBuffer#

A batched circular buffer that returns stochastically delayed values.

Maintains a rolling history of (nworld, *shape) tensors on the GPU and serves values from T - lag timesteps ago, where lag is sampled independently per environment from [min_lag, max_lag].

The buffer is lazily initialized on the first append() or fill() call so it adapts to arbitrary tensor shapes without upfront configuration.

Example usage:

buf = DelayBuffer(min_lag=1, max_lag=3, nworld=64, device="cuda:0")
buf.append(obs)          # obs shape: (64, 6)
delayed = buf.read()     # delayed shape: (64, 6) — each env gets a
                         # different historical frame

Parameters:

min_lag – Minimum delay in substeps (inclusive). 0 means no delay.
max_lag – Maximum delay in substeps (inclusive). Must be >= min_lag.
nworld – Number of parallel environments.
device – Torch device for all tensor allocations.

property min_lag: int#

property max_lag: int#

property empty: bool#: Whether no frame was appended or filled yet.

reset(env_mask: torch.Tensor = None) → None#

Clears history and resamples lags for the given environments.

Parameters:: env_mask – Boolean environment mask to reset. If None, resets all.

append(x: torch.Tensor) → None#

Pushes a new frame into the buffer.

Parameters:: x – Tensor of shape (nworld, *shape) containing the new frame for all environments.

read() → torch.Tensor#

Returns the delayed frame for each environment.

Returns:: Tensor of shape (nworld, *shape).
Raises:: RuntimeError – If no frames have been appended or filled yet.

fill(x: torch.Tensor, env_mask: torch.Tensor) → None#: Fills every slot of env_mask envs (None: all) with x.

class FreeDevice#

Physics-driven 6-DOF free-floating body (cube, ball, etc).

property name: str#: Returns the instance name of this device.

configure_spec(spec: mujoco.MjSpec) → mujoco.MjSpec#: Validate/configure the spec for this device.

initialize(simulator: Simulator) → None#: Resolve body/joint IDs from compiled model.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: No-op. State is controlled by resetter via set_pos/set_quat/set_vel.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Spec for configured substep measurements.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#: Return configured substep measurements.

set_pos(env_mask: torch.Tensor, pos: torch.Tensor) → None#: Set position for specified environments.

set_quat(env_mask: torch.Tensor, quat: torch.Tensor) → None#: Set quaternion for specified environments.

set_vel(env_mask: torch.Tensor, vel: torch.Tensor) → None#: Set 6-DOF velocity (linear + angular) for specified environments.

class MocapDevice#

Fixed-base kinematic marker controlled via mocap (target position).

A puppet entity: exposes set_pos/set_quat for external control. The resetter owns sampling logic and calls these methods.

configure_spec(spec: mujoco.MjSpec, *, geom_group: int | None = None) → mujoco.MjSpec#: Configures MjSpec sensors and optional geom group for the mocap body.

property name: str#: Returns the instance name of this device.

initialize(simulator: Simulator) → None#: Resolve mocap ID from compiled model.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: No-op. Mocap position is controlled by resetter.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Spec for mocap_pos and mocap_quat measurements.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#: Return mocap position and quaternion.

set_pos(env_mask: torch.Tensor, pos: torch.Tensor) → None#: Set mocap position for specified environments.

set_quat(env_mask: torch.Tensor, quat: torch.Tensor) → None#: Set mocap quaternion for specified environments.

class MeasurementDevice#

Entity that wraps callable functions for measurements.

property name: str#: Returns the instance name of this device.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the device state in the simulation.

initialize(simulator: Simulator) → None#: Move tensor params to simulator device and infer missing output shapes.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep measurements spec for this device.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#: Compute measurements by calling the provided functions.

class FunctionalMetricProvider#

Computes per-step metrics and accumulates episodic reductions.

update_metrics(features: Mapping[str, torch.Tensor]) → None#

Evaluates step metrics and updates internal episode accumulators.

Modifies internal state (_step_counts, _mean_sums, _sums, _maxs, _lasts); should only be called once per control step by TaskLayer.update_metrics (via Environment.step).

get_episode_metrics(env_mask: torch.Tensor) → dict[str, torch.Tensor]#

Returns 0-D mean reduced episode metrics over the masked worlds.

Read-only; does not modify internal state.

reset(env_mask: torch.Tensor = None) → None#: Mutates state: clears episode accumulators for the masked worlds.

class FunctionalObservationAdapter#

Wraps a set of ObsGroups as an ObservationSpaceAdapter.

Each observation Obs is evaluated and post-processed in order: func -> bias -> noise -> scale -> offset -> clip -> history

bias — optional noise.Noise sampled once per episode on init/reset (only applied when ObsGroup.enable_noise=True).
noise — optional noise.Noise sampled on every step (only applied when ObsGroup.enable_noise=True).
scale — optional multiplicative factor.
offset — optional additive offset.
clip — optional {"min": ..., "max": ...} clamp bounds.
history_length — optional int number of past steps to stack (oldest to newest: [t - H, ..., t - 1, t]).

reset(env_mask: torch.Tensor = None) → None#

Reset history and episodic bias buffers for the given environment mask.

Zeros the history buffer, resamples per-episode biases for masked worlds, and resets the steps-since-reset counter. On the next observations_from_features call, envs with the first step after reset get their entire history filled with the current observation.

task_features_keys() → set[str]#

observation_spec() → Mapping[str, dm_env.specs.Array]#

observations_from_features(features: Mapping[str, torch.Tensor]) → dict[str, torch.Tensor]#: Computes observations from features (func -> bias -> noise -> scale -> offset -> clip -> history).

class FunctionalEnvironmentReset#

EnvironmentReset with term-based startup/reset.

Terms are filtered by params[TYPE_KEY]: - “startup” -> called in on_startup() - “reset” -> called in on_reset()

Term functions receive device_coordinator as first arg.

For on_step(), subclass and override.

Example:

def randomize_friction(dc, range):
  model = dc.simulation_device.model
  ...

reset = FunctionalEnvironmentReset(
    terms=[
        Term(randomize_friction, type="startup", range=(0.3, 1.5)),
        Term(reset_robot, type="reset", entity="robot"),
    ],
)

set_env(env: environment.Environment) → None#: Called by Environment.__init__.

on_startup() → None#: Call startup terms.

on_reset(config: ResetOptions, env_mask: torch.Tensor) → None#: Call reset terms.

on_step() → None#: Override in subclass for timer-based resampling.

class FunctionalRewardProvider#

Wraps callable functions as a reward provider.

initialize(device: torch.device, dtype: torch.dtype, dt: float = 1.0) → None#: Initialize weight tensor on device. Call before compute_reward.

property name: str#

required_features_keys() → Set[str]#

reward_spec() → RewardSpecDict#

compute_reward(required_features: Mapping[str, torch.Tensor]) → RewardDict#: Compute weighted sum of reward terms, scaled by dt at the end.

reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of the reward provider.

class PassthroughRewardProvider#

Passes a single reward term through without scaling or weighting.

Returns the term’s raw output directly from compute_reward as a torch.Tensor.

initialize(device: torch.device, dtype: torch.dtype, dt: float = 1.0) → None#: Store dtype for reward_spec. Device and dt are unused.

property name: str#

required_features_keys() → Set[str]#

reward_spec() → dm_env.specs.Array#: Returns a scalar reward spec.

compute_reward(required_features: Mapping[str, torch.Tensor]) → torch.Tensor#: Returns the term’s reward value directly.

reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of the reward provider.

class FunctionalTerminationChecker#

Wraps callable functions as a termination checker.

initialize(nworld: int, device: torch.device) → None#: Initialize result buffer. Call before check_termination.

property name: str#

required_features_keys() → Set[str]#

check_termination(features: Mapping[str, torch.Tensor]) → torch.Tensor#: Run all check functions and combine results.

reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of the termination checker.

class InsideSiteDevice#

Containment sensor entity wrapping mjSENS_INSIDESITE.

property name: str#: Returns the instance name of this device.

configure_spec(spec: mujoco.MjSpec) → mujoco.MjSpec#

Add an mjSENS_INSIDESITE sensor to the spec.

The sensor evaluates to 1.0 when the body’s centre of mass lies inside the site’s axis-aligned bounding volume, and 0.0 otherwise.

obj (body) — the body whose CoM is tested for containment. ref (site) — the site that defines the bounding volume.

See https://mujoco.readthedocs.io/en/stable/XMLreference.html#sensor for full mjSENS_INSIDESITE documentation.

initialize(simulator: Simulator) → None#: Initializes the device with the compiled simulator.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the device state in the simulation.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep measurements spec for this device.

get_measurements(simulator: Simulator, *, force: bool = False) → Mapping[str, torch.Tensor]#

Returns the measurements for the current state.

Should not advance any state: DeviceLayer can call this method any number of times and it should return valid measurements at any point in the step lifecycle.

Parameters:: simulator – The simulator.
Returns:: The measurements, batched over worlds. They may be views of simulator or device state, so callers must not modify them in place, and must copy them to keep them across simulator updates.

randomize_actuator_gains(model: TorchWrapper[mujoco_warp.Model], spec: mujoco.MjSpec, device: torch.device, *, pattern: str | tuple[str, ...], kp_noise: Noise | None = None, kd_noise: Noise | None = None, env_mask: torch.Tensor = None) → None#: Randomize kp/kd for affine-biastype actuators, maintaining biasprm[1]=-kp.

randomize_field(model: TorchWrapper[mujoco_warp.Model], spec: mujoco.MjSpec, device: torch.device, *, field: str, pattern: str | tuple[str, ...], noise: Noise | dict[int, Noise], env_mask: torch.Tensor = None, custom_post_processors: None | Callable[[...], None] | Sequence[Callable[[...], None]] = None) → None#

Randomize model fields using pattern-based binding.

Parameters:

model – torch_wrapper.TorchWrapper[mjw.Model] from sim_device.model.
spec – MjSpec for pattern matching.
device – Torch device.
field – Full field name (e.g., “geom_friction”, “body_mass”).
pattern – Name pattern(s) to match entities (regex supported).
noise – Noise configuration (e.g., Uniform, Gaussian) or a dict of noise configurations per axis.
env_mask – Optional boolean mask of environments to randomize.
custom_post_processors – Optional custom post-processors to apply to the field after randomization. If None, default post-processors are applied.

randomize_quat(model: TorchWrapper, spec: mujoco.MjSpec, device: torch.device, *, field: str, pattern: str | tuple[str, ...], max_angle_degrees: float = 3.0, env_mask: torch.Tensor = None) → None#

Randomize quaternion orientation.

Applies a small random rotation perturbation to the matched quaternions.

Parameters:

model – TorchWrapper[mjw.Model].
spec – mjspec for pattern matching.
device – torch device.
field – quaternion field name (e.g., “cam_quat”, “body_quat”, “geom_quat”).
pattern – name pattern(s) to match entities.
max_angle_degrees – maximum rotation perturbation range in degrees.
env_mask – optional boolean mask of environments to randomize.

class MujocoCameraDevice#

Entity that renders the scene using MuJoCo rendering.

property name: str#: Instance name.

property rendering_configuration: Mapping[str, Any]#: The rendering configuration.

initialize(simulator: Simulator) → None#: Initializes the entity with compiled model and data.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#

Resets the processor.

Parameters:

simulator – The simulator instance.
env_mask – Boolean mask of the environments to reset. If None, all environments are reset.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Spec for the produced substep outputs.

on_substep(simulator: Simulator) → None#: Records the delayed qpos and renders on render ticks.

on_begin_stepping(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Records the reset qpos and renders the first frame after a reset.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#: Returns the latest rendered frames, without rendering.

render_now(simulator: Simulator, render_mode: RenderMode | None = None) → dict[str, torch.Tensor]#

Renders the camera immediately without affecting device measurements.

Side-channel call for Impl.CPP single-world environments that renders the current physics state (or the state from delay_ms ago when a delay is configured) without modifying cached substep measurements, the render_rate_hz schedule, or the delay history buffer.

Parameters:

simulator – The simulator instance.
render_mode – Optional subset of enabled RenderMode flags to render. None renders all modes configured on this device.

Returns:

A dictionary containing the freshly rendered images for the requested modes.

Raises:

RuntimeError – If the device has not been initialized.
ValueError – If self._impl != types.Impl.CPP, data.nworld is greater than 1, or render_mode is empty or includes modes not enabled on this device.

substep_commands_spec() → Mapping[str, dm_env.specs.Array]#: No commands are consumed.

set_substep_commands(simulator: Simulator, consumed_substep_commands: Mapping[str, torch.Tensor]) → None#: No commands are consumed.

close() → None#: Closes and releases the underlying MuJoCo renderer, if initialized.

class RenderMode#

Render modes for the various camera entities.

These flags can be combined to render multiple outputs.

RGB = 1#

DEPTH = 2#

SEGMENTATION = 4#

class MultiAsset#

An asset contains a spec with a single body with one or more geoms.

Assets are independent: each is compiled on its own, with its own compiler settings, default classes and file paths, so assets can come from anywhere.

Note

spec should hold exactly one body under its worldbody.
Geoms can be meshes or primitives (sphere, capsule, ellipsoid, cylinder, box).
Inertia comes from the compiled asset.

spec#

MjSpec with one body that may contain multiple geoms.

Type:: mujoco.MjSpec

weight#

Relative weight to sample this asset; must be > 0.

Type:: float

class MultiAssetDevice#

Gives one body a different asset in each world.

property name: str#: Returns the instance name of this device.

configure_spec(spec: mujoco.MjSpec) → mujoco.MjSpec#: Replaces the body’s geoms and copies every asset’s meshes into spec.

initialize(simulator: Simulator) → None#: Assigns an asset to every world and writes its values.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Draws a new asset for every reset world if resample_on_reset.

class DeltaActions#

Task-level delta action integrator: pos_t = clip(pos_{t-1} + action * scale, min, max).

A stateful CommandsProcessor that integrates incoming delta actions starting from a configured offset (e.g. home/hover pose) and bounds them within clip limits.

Example:

DeltaActions(
    input_key="actions",
    output_key="robot/ctrl",
    scale=[0.05, 0.05, 0.05, 0.05, 0.05, 0.10],
    offset=[0.0, -0.26, 0.25, 1.13, 0.08, 0.80],
    action_clip=(-1.0, 1.0),
    smoothing=0.5,
    clip=([-3.14, ...], [3.14, ...]),
    nworld=64,
    xpu_device="cuda:0",
)

input_key — command key to consume (default "actions").
output_key — required, the command key to produce.
scale — optional (default 1.0), scalar or per-dimension sequence/tensor.
offset — optional (default 0.0), initial pose vector or scalar.
action_clip — optional, clamp delta action before scaling (tuple or scalar).
smoothing — optional (default 0.0), exponential smoothing factor in [0, 1).
clip — optional (min, max) tuple of bounds (scalar, sequence, tensor, or None for one-sided clipping).

property name: str#

consumed_commands_spec() → Mapping[str, gdm_robotics.interfaces.types.AnyArraySpec]#

produced_commands_keys() → set[str]#

process_commands(consumed_commands: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#: Integrates delta action: current_pos += delta * scale, then clips.

reset(env_mask: torch.Tensor = None) → None#: Reset integrated positions to initial offsets for given envs.

class StochasticSubstepCommands#

Models device command processing in order: delay -> bias -> noise.

Example:

StochasticSubstepCommands(
    keys="robot/ctrl",
    delay=(0.005, 0.01),
    bias=noise.uniform(-0.01, 0.01),
    noise=noise.uniform(-0.002, 0.002),
    nworld=128,
    xpu_device="cuda:0",
    simulation_timestep=0.002,  # injected by the builder
)

Supported modifiers (applied in order delay -> bias -> noise):

delay — optional (min_seconds, max_seconds) tuple converted to substep lags using simulation_timestep.
bias — optional noise.Noise sampled once per episode on reset.
noise — optional noise.Noise sampled on every substep.

property name: str#: Returns a unique string identifier for this object.

reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of this processor.

produced_substep_commands_keys() → set[str]#: Keys of the substep commands produced by this processor.

consumed_substep_commands_spec() → Mapping[str, dm_env.specs.Array]#: Spec of the substep commands consumed by this processor.

process_substep_commands(simulator: Simulator, substep_commands: MutableMapping[str, torch.Tensor]) → None#

Processes substep commands by modifying the dictionary in place.

Parameters:

simulator – The simulator bound to the DeviceLayer.
substep_commands – mutable dictionary of substep commands. Processors can insert, update, or remove command keys in place.

class StochasticSubstepMeasurements#

Models device measurement processing in order: bias -> noise -> delay.

Example:

StochasticSubstepMeasurements(
    keys=("arm_state/joint_pos", "arm_state/joint_vel"),
    bias=noise.uniform(-0.01, 0.01),
    noise=noise.uniform(-0.002, 0.002),
    delay=(0.005, 0.01),
)

Supported modifiers (applied in order bias -> noise -> delay):

bias — optional noise.Noise sampled once per episode on reset.
noise — optional noise.Noise sampled on every substep.
delay — optional (min_seconds, max_seconds) tuple converted to substep lags using simulation_timestep.

For asymmetric actor-critic setups, prefer to apply noise in the observation terms for the policy only.

property name: str#: Returns a unique string identifier for this object.

reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of this processor.

produced_substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Spec of the substep measurements produced by this processor.

consumed_substep_measurements_keys() → set[str]#: Keys of the substep measurements consumed by this processor.

on_substep(simulator: Simulator, substep_measurements: Mapping[str, torch.Tensor]) → None#: Delays the biased and noisy readings, or samples their noise.

on_begin_stepping(simulator: Simulator, measurements: Mapping[str, torch.Tensor], env_mask: torch.Tensor = None) → None#

Fills the delay buffers and noise samples of env_mask worlds.

The noise is sampled for all worlds and only kept for env_mask worlds, which avoids a host-device sync when no world is reset.

Parameters:

simulator – The simulator, holding the final reset state.
measurements – The measurements of the reset state.
env_mask – Boolean mask of the reset worlds. None: all worlds.

process_measurements(simulator: Simulator, measurements: MutableMapping[str, torch.Tensor]) → None#

Replaces readings by their delayed or biased and noisy values.

Keys without delay apply the bias and latest noise sample to the current reading. Readings pass through until on_substep() or on_begin_stepping() provides their delayed frame, bias, or noise.

Parameters:

simulator – The simulator.
measurements – The measurements, processed in place.

class TransformActions#

Task-level action transform: scale -> offset -> clip.

A CommandsProcessor that transforms raw policy actions before they reach the Device Layer. Pair with ActionHistory (which handles raw-action bookkeeping and feature production) in the task config.

Example:

TransformActions(
    input_key="actions",
    output_key="robot/ctrl",
    scale=[0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.044],
    offset=[0.0, -1.1, 1.0, 0.0, 0.5, 0.0, 0.022],
    clip=(-3.14, 3.14),
)

input_key — command key to consume (default "actions").
output_key — required, the command key to produce.
scale — optional (default 1.0), scalar or per-dimension sequence/tensor.
offset — optional, additive offset after scaling.
clip — optional (min, max) tuple of bounds (scalar, sequence, tensor, or None for one-sided clipping) applied after offset.

property name: str#

consumed_commands_spec() → Mapping[str, gdm_robotics.interfaces.types.AnyArraySpec]#

produced_commands_keys() → set[str]#

process_commands(consumed_commands: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#: Apply scale -> offset -> clip.

reset(env_mask=None) → None#: No state to reset.

class CompositeResetter#

Orchestrates multiple child resetters.

set_env(env: Environment) → None#: Called by Environment.__init__.

on_startup() → None#: Called once after set_env(). Domain randomization.

on_reset(config: ResetOptions, env_mask: torch.Tensor) → None#: User-implemented: called in do_reset.

on_step() → None#: Called each step. Mid-episode resampling, curriculum.

check_termination() → torch.Tensor#: Combine termination: any TERMINATE wins.

class DomainRandomizationResetter#

Model randomization at startup and/or reset using Term functions.

on_startup() → None#: Apply each randomization term to the model.

on_reset(config: gdm_robotics.interfaces.environment.Options, env_mask: torch.Tensor) → None#: User-implemented: called in do_reset.

class NanTerminationResetter#

Terminates worlds whose simulator physics state is non-finite.

property name: str#

check_termination() → torch.Tensor#: Checks for NaN or Inf in qpos, qvel, qacc, or qacc_warmstart.

class TerminationResetter#

Sim-only termination checks using Term functions.

on_reset(config: ResetOptions, env_mask: torch.Tensor) → None#: User-implemented: called in do_reset.

check_termination() → torch.Tensor#: Check all termination conditions, any True triggers termination.

class SiteDevice#

Read-only entity for extracting site position and orientation.

property name: str#: Returns the instance name of this device.

configure_spec(spec: mujoco.MjSpec) → mujoco.MjSpec#: Validate/configure the spec for this device.

initialize(simulator: Simulator) → None#: Resolve site ID from compiled model.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: No-op. Sites are read-only.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Spec for site xpos and xmat measurements.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#: Return site xpos and xmat.

class SpecHelper#

Helper to manage and prefix specification names for environment measurements/commands.

property spec_dict: Mapping[str, dm_env.specs.Array]#: Returns a mapping of full prefixed measurement names to array specs.

property names_dataclass: SpecTemplate#: Returns the SpecTemplate instance with fields containing the full prefixed measurement names.

make_dict_of_arrays(**kwargs: torch.Tensor) → Mapping[str, torch.Tensor]#: Returns a mapping of full prefixed measurement names to array values.

class TactileDevice#

Attaches a mjSENS_TACTILE sensor to a geom and mesh and reads depth.

Exposes f"{name}/depth" with shape (nworld, ntaxel) containing raw penetration depth in meters (typically ~1e-3 m). Downstream flat_scene.Obs terms should configure scale=... (e.g. 1 / max_depth) for RL policy inputs and Studio Obs 2D visualization (which clips float maps to [0, 1]).

property name: str#: Returns the instance name of this device.

configure_spec(spec: mujoco.MjSpec, *, geom: str, mesh: str = '', cutoff: float = 0.0) → mujoco.MjSpec#

Attaches a mjSENS_TACTILE sensor to geom and mesh in spec.

Parameters:

spec – The MjSpec being configured.
geom – Reference geom name in spec whose frame positions the taxels.
mesh – Mesh name in spec whose vertices define the taxels. Defaults to geom’s mesh when geom is a mesh geom.
cutoff – Optional non-negative sensor cutoff value.

Returns:

The configured MjSpec.

initialize(simulator: Simulator) → None#: Resolves sensor address, taxel count, and optional Laplacian inverse.

reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the device state in the simulation.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Returns the unbatched measurement spec ({} when dark=True).

get_measurements(simulator: Simulator, *, force: bool = False) → Mapping[str, torch.Tensor]#: Returns a batched (nworld, ntaxel) tensor of tactile depths.

compute_laplacian_inv(verts: numpy.ndarray, faces: numpy.ndarray) → numpy.ndarray#: Precomputes the (nvert, nvert) inverse cotangent Laplacian.

class TimeoutTerminationChecker#

Truncates when the episode time exceeds the specified timeout.

name() → str#

property timeout: float#: Returns the episode timeout in seconds.

check_termination(required_features: Mapping[str, torch.Tensor]) → torch.Tensor#: Checks if the episode should terminate.

required_features_keys() → set[str]#

class ZeroRewardProvider#

Reward provider which provides a zero reward.

name() → str#

compute_reward(required_features: Mapping[str, torch.Tensor]) → torch.Tensor#: Returns a zero reward.

reward_spec() → tree.Structure.dm_env.specs.Array#: Returns the spec for a constant zero reward.

required_features_keys() → set[str]#

Returns empty set.

There are no feature keys that are required to compute the reward.

reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of the reward provider.

class Clock#

Discrete physics and environment clock.

sim_dt#

Physics simulation timestep (seconds).

Type:: float

control_dt#

Environment control timestep (seconds).

Type:: float

n_substeps#

Number of physics substeps per control step.

Type:: int

nworld#

Number of parallel environments / worlds.

Type:: int

xpu_device#

Device on which episode counters are allocated.

Type:: torch.device

global_step#

Monotonically increasing count of completed control steps.

Type:: int

global_substep#

Monotonically increasing count of physics substeps since simulation start.

Type:: int

episode_step#

Tensor of shape (nworld,) containing control steps since last reset for each environment.

Type:: torch.Tensor

episode_substep#

Tensor of shape (nworld,) containing physics substeps since last reset for each environment.

Type:: torch.Tensor

property global_time: float#: Current continuous simulation time in seconds.

property episode_time: torch.Tensor#: Current continuous episode time in seconds for each environment.

tick_substep() → None#: Advances the physics substep counters by 1.

tick_step() → None#: Advances the control step counters by 1.

reset(env_mask: torch.Tensor = None) → None#

Resets episode step and substep counters for specified environments.

Parameters:: env_mask – Boolean tensor of shape (nworld,) indicating which environments to reset. If None, resets all environments.

class Device#

A REAF Device representing a simulation participant.

configure_spec(spec: mujoco.MjSpec) → mujoco.MjSpec#: Validate/configure the spec for this device.

initialize(simulator: Simulator) → None#: Initializes the device with the compiled simulator.

abstract property name: str#: Returns the instance name of this device.

abstractmethod reset(simulator: Simulator, env_mask: torch.Tensor = None) → None#: Resets the device state in the simulation.

substep_commands_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep commands spec for this device.

substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Returns the substep measurements spec for this device.

set_substep_commands(simulator: Simulator, consumed_substep_commands: Mapping[str, torch.Tensor]) → None#: Sets the substep commands for this device.

on_substep(simulator: Simulator) → None#

Advances per-substep state after every physics substep.

E.g. pushes the current reading into a substep history buffer.

Parameters:: simulator – The simulator, after the physics substep.

on_step(simulator: Simulator) → None#

Advances per-control-step state on the final substep.

Runs on the last physics substep right after on_substep() and before get_measurements() and substep processors, holding the prior step’s value across earlier substeps (zero-order hold). Does not run after resets: on_begin_stepping() sets this state for reset worlds.

Parameters:: simulator – The simulator, after the last physics substep.

on_begin_stepping(simulator: Simulator, env_mask: torch.Tensor = None) → None#

Records the reset state of the env_mask worlds.

reset() runs before the environment resetter, Simulator.set_const(), and Simulator.forward(), so simulator.data does not hold the reset state yet. This hook runs after them and before the reset state is read, e.g. to fill history buffers with the reset reading. Worlds outside env_mask must keep their state.

Parameters:

simulator – The simulator, holding the final reset state.
env_mask – Boolean mask of the reset worlds. None: all worlds.

get_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#

Returns the measurements for the current state.

Should not advance any state: DeviceLayer can call this method any number of times and it should return valid measurements at any point in the step lifecycle.

Parameters:: simulator – The simulator.
Returns:: The measurements, batched over worlds. They may be views of simulator or device state, so callers must not modify them in place, and must copy them to keep them across simulator updates.

get_substep_measurements(simulator: Simulator) → Mapping[str, torch.Tensor]#

Deprecated: implement on_substep(), on_step(), and get_measurements().

DeviceLayer calls overrides after every physics substep and after resets.

Parameters:: simulator – The simulator.
Returns:: The measurements, batched over worlds.

class DeviceLayer#

Simulation device layer.

Coordinates:

The simulator (MjSpec -> Model & Data, physics stepping).
All Devices.
All substep command and measurement processors.
The substep loop.

property xpu_device: str#

property nworld: int#

property devices: Sequence[Device]#

property sim: Simulator#

forward() → None#

commands_spec() → Mapping[str, gdm_robotics.interfaces.types.AnyArraySpec]#

measurements_spec() → Mapping[str, dm_env.specs.Array]#

begin_stepping(env_mask: torch.Tensor = None) → Mapping[str, torch.Tensor]#: Records the reset state, then reads, updating only env_mask rows.

end_stepping(env_mask: torch.Tensor = None) → None#

step(commands: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#

Runs one control step and returns the measurements of all worlds.

Parameters:: commands – The commands for this control step.
Returns:: The measurements of all worlds. The same tensors are returned and overwritten in place by every step(), begin_stepping(), and update_measurements(), so callers that keep measurements across these calls (e.g. for a history) must .clone() them.

update_measurements() → None#: Reads all devices and processes their measurements, without hooks.

get_measurements() → Mapping[str, torch.Tensor]#

Returns the latest measurements of all worlds.

The tensors are overwritten in place by every step(), begin_stepping(), and update_measurements(), so callers that keep them across these calls must .clone() them.

property control_timestep: float#

property clock: Clock#

has_device(name: str) → bool#

get_device(name: str) → Device#

reset(env_mask: torch.Tensor = None) → None#

class DiscountProvider#

Computes the discount.

compute_discount(required_features: Mapping[str, torch.Tensor], termination_state: torch.Tensor) → torch.Tensor#: Computes the discount.

reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of the discount provider.

class AutoResetMode#

Controls how Environment.step() handles episode auto resets.

SYNC#: Default. At the end of step(), checks on the host whether any world’s episode ended on the previous step (a host-device sync), and only then resets those worlds. This avoids the expensive mj_forward calls that default resetters (e.g. EnvironmentReset) execute when no world needs resetting.

NO_SYNC#: At the end of every step(), resets the worlds whose episode ended on the previous step, without checking on the host first. This avoids a host-device sync, but may invoke an expensive mj_forward on every step (e.g. the default EnvironmentReset implementation).

DISABLED#: Disables automatic resetting inside Environment.step(). The caller is responsible for calling reset_with_options(options=..., env_mask=...) explicitly. Gym/VecEnv style wrappers should use this option to avoid the default auto-reset and CUDA syncs in Environment.step().

class EndOfEpisodeHandler#

Handler called after the last episode step.

on_end_of_episode_stepping(final_timestep: dm_env.TimeStep, env_mask: torch.Tensor) → None#

class Environment#

The Robotics Environment Authoring Framework (REAF) Environment class.

TODO(b/513593452, btaba): Bake in TaskState protocol and get_state()/set_state() public APIs on Environment and TaskLayer to avoid ad-hoc goal state poking.

property xpu_device: str#: Returns the hardware device used by the environment.

property nworld: int#: Returns the number of worlds running in simulation.

property auto_reset_mode: AutoResetMode#: Returns the automatic reset mode for step().

default_reset_options() → ResetOptions#

action_spec() → Any#: Defines the actions that should be provided to step.

timestep_spec() → gdm_robotics.interfaces.types.TimeStepSpec#: Returns the spec associated to the returned TimeStep.

reset_with_options(*, options: ResetOptions, env_mask: torch.Tensor = None, is_last_timestep: torch.Tensor = None) → dm_env.TimeStep#: Starts a new sequence and returns the first TimeStep.

step(action: gdm_robotics.interfaces.types.ActionType) → dm_env.TimeStep#: Updates the environment according to action and returns a TimeStep.

property device_layer: Any#

property task_layer: TaskLayer#

property environment_reset: EnvironmentReset#

add_logger(logger: reaf.core.logger.Logger) → None#

remove_logger(logger: reaf.core.logger.Logger) → None#

class EnvironmentReset#

Reset hooks for environment lifecycle with privileged sim access.

Hooks:

on_startup(): Called once at startup (e.g. domain randomization)
on_reset(): Called on every reset (e.g. initialize new episode targets)
on_step(): Called on every step (e.g. goal resampling, curriculum)
check_termination(): Triggers termination using privileged sim state (e.g. NaN detection, joint limits) — unlike TaskLayer terminators, this has direct access to model/data. Called on every step.

The resetter receives full Environment access via set_env(), called by Environment.__init__.

set_env(env: Any) → None#: Called by Environment.__init__.

do_reset(config: ResetOptions, env_mask: torch.Tensor) → None#

Restore model defaults, reset the device layer, and call on_reset().

Model fields written through sim.model are tracked by the model TorchWrapper. The Environment then decides if it should call Simulator.set_const(), followed by a hard-coded call to forward().

Parameters:

config – The reset options.
env_mask – Boolean mask of the worlds to reset (None resets all).

Raises:

RuntimeError – If set_env() was not called before do_reset().

on_startup() → None#: Called once after set_env(). Domain randomization.

on_reset(config: ResetOptions, env_mask: torch.Tensor) → None#: User-implemented: called in do_reset.

on_step() → None#: Called each step. Mid-episode resampling, curriculum.

check_termination() → torch.Tensor#: Check for termination using privileged sim state.

default_reset_configuration() → ResetOptions#: Returns the default reset configuration.

class FeatureProducer#

Produces additional features to be exposed by the task layer.

This is the torch version of the base FeaturesProducer class.

abstractmethod produce_features(required_features: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#: Produces additional features for the environment.

reset(env_mask: torch.Tensor = None) → None#

Resets the internal state of the feature producer.

Parameters:: env_mask – Optional boolean environment mask to reset. If None, resets all environments.

class MetricProvider#

Computes episode metrics from per-step task features.

abstractmethod update_metrics(features: Mapping[str, torch.Tensor]) → None#

Evaluates step metrics and updates internal episode accumulators.

Modifies internal state; should only be called once per control step by TaskLayer.update_metrics (via Environment.step).

reset(env_mask: torch.Tensor = None) → None#: Mutates state: clears episode accumulators for the masked worlds.

abstractmethod get_episode_metrics(env_mask: torch.Tensor) → dict[str, torch.Tensor]#

Returns 0-D mean reduced episode metrics over masked worlds.

Read-only; does not modify internal state.

get_spec_by_obj_type(spec: mujoco.MjSpec, obj_type: mujoco.mjtObj, name: str) → mujoco.MjStruct#: Maps obj_type to the corresponding MjSpec accessor.

match_names(patterns: Sequence[str], names: Sequence[str]) → list[str]#

Returns the names that fully match any regex in patterns.

Parameters:

patterns – Name regexes.
names – Candidate names (e.g. joints or actuators).

Returns:

The matching names, in names order.

resolve_names(spec: mujoco.MjSpec, objtype: mujoco.mjtObj, pattern: str | tuple[str, ...]) → list[str]#

Resolve names from spec via pattern matching.

Parameters:

spec – MjSpec to search.
objtype – mujoco.mjtObj (mjOBJ_GEOM, mjOBJ_BODY, mjOBJ_XBODY, mjOBJ_ACTUATOR).
pattern – Name(s) to match (exact, glob, or regex). Tuple for multiple.

Returns:

List of matched names.

Raises:

ValueError – If no names match the pattern.

resolve_values(patterns: Mapping[str, float], names: Sequence[str], default: float = 0.0) → list[float]#

Resolves a {regex: value} mapping to a per-name value list.

Each regex must fully match a name. A name may match at most one pattern; names matching none get default.

Parameters:

patterns – Map from name regex to the value for matching names.
names – Names (e.g. joints or actuators), in the order of the returned list.
default – Value for names that match no pattern.

Returns:

One value per entry in names.

Raises:

ValueError – If a name matches more than one pattern.

soft_joint_limits(model: mujoco.MjModel, joint_names: Sequence[str], factor: float) → numpy.ndarray#

Returns joint ranges scaled by factor about their midpoints.

Parameters:

model – Compiled model containing the joints.
joint_names – Joints to read, in the order of the returned rows.
factor – Fraction of each joint’s full range to keep (e.g. 0.9).

Returns:

Array of shape (len(joint_names), 2) holding [low, high] per joint.

class ObservationSpaceAdapter#

reaf-sim specific ObservationSpaceAdapter interface.

Overrides the reset method to provide torch-specific type annotations for environment indices.

observations_from_features(required_features: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#: Computes the observation from the required features.

reset(env_mask: torch.Tensor = None) → None#

Reset any internal state for the given environment mask.

Called by the environment on episode reset. Default is a no-op. Override this method to support stateful buffers (e.g. history).

Parameters:: env_mask – Boolean environment mask to reset. If None, resets all.

class TorchDeviceLayer#

Protocol defining the reaf-sim simulation DeviceLayer interface.

forward() → None#: Advances/recomputes simulation kinematic state to match current positions.

property nworld: int#: Number of parallel worlds / batch dimension.

property xpu_device: str#

0’, ‘cpu’).

Type:: Device hosting tensors (e.g. ‘cuda

property sim: Simulator | None#: Underlying simulator instance, or None for hardware device layers.

property control_timestep: float#: Control timestep in seconds.

property clock: Clock | None#: Simulation clock, or None for hardware device layers.

measurements_spec() → Mapping[str, dm_env.specs.Array]#: Returns the batched measurements specification.

commands_spec() → Mapping[str, gdm_robotics.interfaces.types.AnyArraySpec]#: Returns the batched commands specification.

begin_stepping(env_mask: torch.Tensor = None) → Mapping[str, torch.Tensor]#: Prepares and returns initial measurements before stepping.

step(commands: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#: Steps the device layer with commands and returns measurements.

end_stepping(env_mask: torch.Tensor = None) → None#: Completes the stepping phase.

reset(env_mask: torch.Tensor = None) → None#: Resets simulator and device states.

get_measurements() → Mapping[str, torch.Tensor]#: Returns the latest measurements dictionary.

get_device(name: str) → Any#: Returns the device with the specified name.

has_device(name: str) → bool#: Returns True if a device with the specified name exists.

class RewardProvider#: Computes the reward.

class Simulator#

MuJoCo / MuJoCo-Warp physics simulator.

Encapsulates:

MuJoCo MjSpec and compiled MjModel.
Raw and Torch-wrapped Model and Data tensors.
Physics stepping, forward kinematics, and state resets.
CUDA graph capture management for Warp backend.

compile() → None#: Compiles MjSpec to MjModel/MjData and allocates backend models.

property clock: Clock#: Simulation clock tracking step and substep counts.

set_clock(clock: Clock) → None#: Sets the simulation clock.

property model: TorchWrapper[mujoco_warp.Model]#: Torch-wrapped model attributes across batch dimensions.

property data: TorchWrapper[mujoco_warp.Data]#: Torch-wrapped dynamic state tensors across batch dimensions.

property warp_model: mujoco_warp.Model#

property warp_data: mujoco_warp.Data#

property mj_spec: mujoco.MjSpec#

property mj_model: mujoco.MjModel#: Underlying compiled CPU MuJoCo model.

property mj_data: mujoco.MjData#: Underlying CPU MuJoCo data buffer.

property impl: Impl#

property nworld: int#: Number of parallel worlds / batch dimension.

property xpu_device: str#

0’, ‘cpu’).

Type:: Device hosting tensors (e.g. ‘cuda

property nvmax: int | None#

forward() → None#: Computes forward kinematics and dynamics.

reset(env_mask: torch.Tensor | None = None) → None#: Resets simulation data for specified environments.

step() → None#: Executes a single low-level physics step.

step_n(n_substeps: int) → None#: Executes n_substeps low-level physics steps.

set_const() → SetConstLevel#

Recomputes derived model constants after primary model fields changed.

Flushes the SetConstLevel tracked by model (see TorchWrapper.flush_set_const) and runs the matching mj_setConst stages for all worlds. Only in-place writes through self.model (TorchWrapper) are tracked; direct modifications to the base models (mj_model, warp_model, ._tensor, or ._array) are not tracked. This is a no-op if no tracked model field was written since the last call. The Warp backend always runs with restore=False, so derived data fields (e.g. xpos, qM) are left at qpos0 for all worlds; callers must call forward() before reading data.

Returns:: The SetConstLevel that was run.

class SubstepCommandsProcessor#

reaf-sim specific substep commands processor.

abstract property name: str#: Returns a unique string identifier for this object.

abstractmethod reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of this processor.

initialize(simulator: Simulator) → None#: Initializes the processor with the compiled model.

abstractmethod produced_substep_commands_keys() → set[str]#: Keys of the substep commands produced by this processor.

abstractmethod consumed_substep_commands_spec() → Mapping[str, dm_env.specs.Array]#: Spec of the substep commands consumed by this processor.

abstractmethod process_substep_commands(simulator: Simulator, substep_commands: MutableMapping[str, torch.Tensor]) → None#

Processes substep commands by modifying the dictionary in place.

Parameters:

simulator – The simulator bound to the DeviceLayer.
substep_commands – mutable dictionary of substep commands. Processors can insert, update, or remove command keys in place.

class SubstepMeasurementsProcessor#

reaf-sim specific substep measurements processor.

abstract property name: str#: Returns a unique string identifier for this object.

abstractmethod reset(env_mask: torch.Tensor = None)#: Resets the internal state of this processor.

initialize(simulator: Simulator) → None#: Initializes the processor with the compiled model.

abstractmethod produced_substep_measurements_spec() → Mapping[str, dm_env.specs.Array]#: Spec of the substep measurements produced by this processor.

abstractmethod consumed_substep_measurements_keys() → set[str]#: Keys of the substep measurements consumed by this processor.

on_substep(simulator: Simulator, substep_measurements: Mapping[str, torch.Tensor]) → None#

Advances per-substep state after every physics substep.

E.g. pushes the substep reading into a delay buffer.

Parameters:

simulator – The simulator, after the physics substep.
substep_measurements – The measurements of the substep, as processed by the preceding processors. They may be views of simulator or device state, so copy them into processor-owned buffers to keep them (e.g. in a history), and do not modify them in place.

on_begin_stepping(simulator: Simulator, measurements: Mapping[str, torch.Tensor], env_mask: torch.Tensor = None) → None#

Records the reset state of the env_mask worlds.

Runs after reset(), once the reset state is final and before it is processed, e.g. to fill delay buffers with the reset reading. Worlds outside env_mask must keep their state.

Parameters:

simulator – The simulator, holding the final reset state.
measurements – The measurements of the reset state, as processed by the preceding processors. They may be views of simulator or device state, so copy them into processor-owned buffers to keep them (e.g. in a history), and do not modify them in place.
env_mask – Boolean mask of the reset worlds. None: all worlds.

process_measurements(simulator: Simulator, measurements: MutableMapping[str, torch.Tensor]) → None#

Transforms measurements in place.

Should not advance any state: DeviceLayer may call this method more than once per substep (e.g. after resets or in update_measurements()).

Parameters:

simulator – The simulator.
measurements – Mutable dictionary of measurements. Processors can insert, update, or remove measurement keys in place, but must not modify the tensors of measurements in place (they may be views of simulator or device state); they may write into tensors they own (e.g. out=).

process_substep_measurements(simulator: Simulator, substep_measurements: MutableMapping[str, torch.Tensor]) → None#

Processes substep measurements by modifying the dictionary in place.

Deprecated: implement on_substep() and process_measurements() instead. DeviceLayer calls overrides after every physics substep and after resets, and emits a DeprecationWarning for them.

Parameters:

simulator – the simulator bound to the DeviceLayer.
substep_measurements – mutable dictionary of substep measurements. Processors can insert, update, or remove measurement keys in place.

class TaskLayer#

Torch-compatible version of the REAF TaskLayer class.

commands_spec(device_layer_commands_spec: Mapping[str, gdm_robotics.interfaces.types.AnyArraySpec] | None = None) → Mapping[str, gdm_robotics.interfaces.types.AnyArraySpec]#: Returns the commands spec exposed by the task layer.

features_spec(device_layer_measurements_spec: Mapping[str, dm_env.specs.Array] | None = None) → dict[str, dm_env.specs.Array]#: Returns the features spec exposed by the task layer.

validate_spec(*, device_layer_commands_spec: Mapping[str, gdm_robotics.interfaces.types.AnyArraySpec] | None = None, device_layer_measurements_spec: Mapping[str, dm_env.specs.Array] | None = None) → None#: Checks that the specs have consistent keys.

reward_spec()#: Returns the reward spec exposed by the task layer.

discount_spec()#: Returns the discount spec exposed by the task layer.

perform_reset(env_mask: torch.Tensor = None) → None#: Reset the internal state of the task layer.

check_for_termination(features: Mapping[str, torch.Tensor]) → torch.Tensor#: Checks for termination.

compute_reward(features: Mapping[str, torch.Tensor]) → torch.Tensor#: Computes the reward given the features.

update_metrics(features: Mapping[str, torch.Tensor]) → None#

Evaluates per-step metrics and updates internal episodic accumulators.

Modifies internal state; should only be called once per control step by Environment.step.

get_episode_metrics(env_mask: torch.Tensor) → dict[str, torch.Tensor]#

Returns 0-D mean reduced episode metrics over the masked worlds.

Read-only; does not modify internal state.

compute_discount(features: Mapping[str, torch.Tensor], terminations: torch.Tensor) → torch.Tensor#: Computes the discount given the features and terminations.

compute_all_features(measurements: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#

Computes all the task layer features given the current measurements.

Records nothing to observers or loggers: see record().

Parameters:: measurements – The device layer measurements.
Returns:: All the task layer features.

record(measurements: Mapping[str, torch.Tensor], features: Mapping[str, torch.Tensor]) → None#

Passes measurements and features to the observers and loggers.

Parameters:

measurements – The device layer measurements.
features – All the task layer features.

compute_final_commands(commands: Mapping[str, torch.Tensor]) → Mapping[str, torch.Tensor]#: Computes final commands by running commands through all processors.

add_logger(logger: reaf.core.logger.Logger) → None#: Adds a logger to the task layer.

remove_logger(logger: reaf.core.logger.Logger) → None#: Removes a logger from the task layer.

get_component_by_name(name: str)#: Retrieves a component by its name.

has_component(name: str) → bool#: Checks if a component with the given name exists.

class TerminationChecker#

Checks if the episode should terminate.

check_termination(required_features: Mapping[str, torch.Tensor]) → torch.Tensor#: Checks if the episode should terminate.

reset(env_mask: torch.Tensor = None) → None#: Resets the internal state of the termination checker.

class JointType#

dof_width() → int#

qpos_width() → int#

class MjField#

name: str#

is_batch: bool#

shape: tuple[str, ...]#

class MjObjType#

Enum for different types of Mujoco objects.

MJW_MODEL = 'MJW_MODEL'#

MJW_DATA = 'MJW_DATA'#

MJ_DATA = 'MJ_DATA'#

MJ_MODEL = 'MJ_MODEL'#

MJ_OPTION = 'MJ_OPTION'#

MJW_OPTION = 'MJW_OPTION'#

OTHER = 'OTHER'#

classmethod from_obj(obj: Any) → MjObjType#

class SetConstLevel#

Determines which mj_setConst levels must run after modifying the model.

Levels are defined as bitmask flags to support independent compilation stages.

class TorchTensor#

Wraps a warp/numpy array to expose it as a torch.Tensor.

If on_write is set, it is called before explicit write operations. This is typically used to determine when to call set_const after modifying Model fields.

Tracked operation on TorchTensor:

Subscript assignment: model.<field>[idx] = val via __setitem__
Bound attribute assignment: bound.set("field", val) (uses __setitem__)

Write operations that bypass the wrapper are NOT tracked:

Direct writes to underlying storage: t._tensor[idx] = val or t._array[idx] = val
In-place ops on views from t[idx], e.g. t[idx] += val
In-place tensor methods: t.copy_(), t.add_(), t.mul_(), t.zero_(), etc.
In-place torch module functions: torch.add_(t, ...), torch.clamp_(t, ...)
Functions writing to out= destinations: torch.add(..., out=t), etc.
Direct mutation of raw Warp / C++ models: sim.warp_model.<field>, sim.mj_model.<field>.

Access to underlying storage is private (_tensor and _array). Do not access or mutate _tensor or _array directly.

property device: torch.device#

class TorchWrapper#

Wraps objects that contain Warp/Numpy arrays and exposes them as tensors.

property obj: T#

property nworld: int#

property pending_set_const: SetConstLevel#: The SetConstLevel recorded since the last flush_set_const().

flush_set_const() → SetConstLevel#: Returns the pending SetConstLevel and resets it to NONE.

clear_defaults() → None#: Clears saved defaults so current model values become the new defaults.

restore_defaults(env_mask: torch.Tensor) → None#: Restores modified model fields to their defaults in-place for env_mask.

bind(specs: Any, /) → BoundModelAccessor#
bind(model: Any, specs: Any, /) → BoundDataAccessor: Binds a Data or Model object.

bind_pattern(spec: mujoco.MjSpec, obj_type: mujoco.mjtObj, pattern: str | tuple[str, ...]) → BoundModelAccessor#

Binds model fields by name pattern.

Parameters:

spec – MjSpec used for name resolution.
obj_type – mujoco.mjtObj (mjOBJ_GEOM, mjOBJ_BODY, etc).
pattern – Name pattern (exact, glob, or regex). Tuple for multiple.

Returns:

BoundModelAccessor for the matched entities.

Raises:

ValueError – If no names match or obj_type unsupported.
TypeError – If called on non-Model object.

register_graph_capture_callback(callback: Callable[[], None]) → None#: Registers a CUDA graph capture callback if TorchWrapper invalidates memory.

class Impl#

Simulation implementation to use.

CPP = 'cpp'#

WARP = 'warp'#

class Term#

Protocol for term objects with function and params.

func: Callable[[...], Any]#

params: dict[str, Any]#

output_shape: tuple[int, ...] | None#

reaf.simulation.torch.flat_scene#

apply_sim_config(spec: mujoco.MjSpec, sim: Sim) → mujoco.MjSpec#: Applies physics and solver parameters from Sim config to an MjSpec.

class Sim#

Simulation parameters.

nworld: int = 1#

xpu_device: str = 'cuda:0'#

impl: Literal['warp', 'cpp'] = 'warp'#

control_timestep: float = 0.02#

simulation_timestep: float = 0.005#

jacobian: mujoco.mjtJacobian = None#

integrator: mujoco.mjtIntegrator = None#

iterations: int | None = None#

ls_iterations: int | None = None#

impratio: float | None = None#

cone: mujoco.mjtCone = None#

solver: mujoco.mjtSolver = None#

gravity: tuple[float, float, float] | None = None#

ccd_iterations: int | None = None#

nconmax: int | None = None#

njmax: int | None = None#

nccdmax: int | None = None#

nvmax: int | None = None#

build_device_layer(env_config: Environment) → DeviceLayer#: Builds only the simulated device layer.

build_env(env_config: Environment) → Environment#: Build a full REAF Environment from config.

build_scene(env_config: Environment) → mujoco.MjSpec#: Builds the full spec by calling all devices spec functions.

build_task_layer(env_config: Environment) → TaskLayer#: Build only the task layer.

make_deployment_config(env_config: Environment, *, observation_groups: Sequence[str]) → Environment#

Strips a training config down to what a deployed policy needs.

Deployment keeps only the MDP boundary the policy was trained against. Everything that exists for the learner rather than for the robot is dropped, and the config is pinned to a single world.

Parameters:

env_config – The training config the policy was trained with.
observation_groups – The observation groups the policy consumes.

Returns:

A deployment config that can be used to test the policy in sim, or splice on top of a real device layer for sim2real deployment.

class Component#

Unified configuration for REAF devices, processors, etc..

component_type: type[Any]#

build_params: dict[str, Any]#

runtime_params: dict[str, Any]#

spec_fn: Callable[[], mujoco.MjSpec] | None = None#

attachment_prefix: str = ''#

attachment_site: str | None = None#

attachment_pos: tuple[float, float, float] | None = None#

attachment_quat: tuple[float, float, float, float] | None = None#

privileged: bool = False#

configure_spec(*args: ~_P_spec, **build_kwargs: ~_P_spec) → Self#: Configures MjSpec build-time arguments on this component.

attach(spec_fn: Callable[[], mujoco.MjSpec], *, prefix: str, site: str | None = None, pos: tuple[float, float, float] | None = None, quat: tuple[float, float, float, float] | None = None) → Self#: Configures how this device’s MjSpec attaches to the base scene.

mark_privileged() → Self#: Marks this component as privileged so it is pruned in deployment configs.

class DeviceLayer#

Configuration for simulation devices.

devices: dict[str, Component]#

command_processors: dict[str, Component]#

substep_measurement_processors: dict[str, Component]#

class Environment#

Top-level environment configuration.

base_scene: mujoco.MjSpec#

device_layer: DeviceLayer#

task_layer: TaskLayer#

action_space_adapter: reaf.core.action_space_adapter.ActionSpaceAdapter = None#

observation_adapter: dict[str, ObsGroup]#

resetters: dict[str, Component]#

sim: Sim#

episode_length_secs: float | None = 20.0#

nan_termination: bool = False#

class Metric#

Metric term evaluated at each step and reduced over episodes.

reduce: Literal['last', 'max', 'mean', 'sum'] = 'mean'#

class Obs#

Observation term with explicit post-processing modifiers and kwargs.

output_shape: tuple[int, ...] | None = None#

bias: Any = None#

noise: Any = None#

scale: Any = None#

offset: Any = None#

clip: dict[str, Any] | None = None#

history_length: int = 0#

dtype: Any = None#

configure(*, bias: Any = None, noise: Any = None, scale: Any = None, offset: Any = None, clip: dict[str, Any] | None = None, history_length: int = 0, dtype: Any = None, output_shape: tuple[int, ...] | None = None) → Self#: Configures observation post-processing modifiers and metadata.

class ObsGroup#

Group of observation terms with shared settings (e.g.

per-episode bias and per-step noise).

terms: dict[str, Obs]#

enable_noise: bool = False#

class Reward#

Reward term with explicit weight and kwargs.

weight: float = 1.0#

class TaskLayer#

Task logic configuration.

rewards: dict[str, Reward]#

reward_provider: Component | None = None#

metrics: dict[str, Metric]#

terminations: dict[str, Term]#

termination_checkers: dict[str, Component]#

command_processors: dict[str, Component]#

feature_producers: dict[str, Component]#

command_feature_producers: dict[str, Component]#

class Term#

Universal term: function + params.

func: Callable[[...], Any]#

params: dict[str, Any]#

output_shape: tuple[int, ...] | None = None#

reaf.simulation.torch.viewer#

launch(environment: gdm_robotics.interfaces.environment.Environment, policy: gdm_robotics.interfaces.policy.Policy.~PolicyDataType = None, simulator: Simulator | None = None, step_duration: float | None = None, reward_plot_keys: Iterable[str] | None = None, loggers: Iterable[Any] | None = None, main_world_id: int = 0, gfx: str = 'web', http_port: int = 0, auto_reset: bool = True, initial_instruction: str | None = None, plugins: Iterable[mujoco.studio.plugin.Plugin] | None = None, hot_reload: str = '') → None#

Launches a viewer for a reaf-sim environment.

Parameters:

environment – The reaf-sim environment to visualize.
policy – An optional policy to execute. If None, use a zero-action policy.
simulator – The physics simulator. If None, infer from environment.
step_duration – Optional desired duration of each step in seconds.
reward_plot_keys – The reward plot legends to show, in order. Nested reward dict keys are joined with ‘/’ and vector entries are indexed, e.g. terms/height or [0]. If None, all rewards are plotted.
loggers – Optional iterable of EpisodicLogger to log episode data.
main_world_id – Index of the world to visualize in the viewer.
gfx – Graphics mode for the viewer (e.g. ‘classic_headless’, ‘web’).
http_port – Web Viewer port (0 picks first free port >= 8080).
auto_reset – If True (default), the environment automatically resets when an episode ends. If False, the viewer pauses instead, allowing inspection of the final state. The user can still manually reset with backspace.
initial_instruction – Optional initial instruction to send to the viewer.
plugins – Optional viewer and sim plugins.
hot_reload – What is reloaded when modified. Options are ‘’, ‘plugins’, ‘all’, ‘Plugin1,Plugin2,…’. Default is ‘’, meaning nothing. See mujoco.studio.hot_reload.