embodichain.lab.gym

Gymnasium-compatible robot learning environment framework.

Extends the Gymnasium API with multi-environment parallel execution, custom observations, and robotic-specific functionality.

Submodules

envs	Environment framework: BaseEnv / EmbodiedEnv class hierarchy, task registration, manager wiring, and the step/reset lifecycle.
utils	Environment utilities: the registration system (register_env, make), Gymnasium integration helpers, miscellaneous utilities, and EnvProfiler for step/reset timing.

Overview

The gym module provides a comprehensive framework for creating robot learning environments. It extends the Gymnasium interface to support multi-environment parallel execution, custom observations, and robotic-specific functionality.

Key Features:

• Multi-Environment Support: Run multiple environment instances in parallel for efficient training

• Gymnasium Integration: Full compatibility with the Gymnasium API and ecosystem

• Robotic Focus: Built-in support for robot control, sensors, and manipulation tasks

• Extensible Architecture: Easy to create custom environments and tasks

• GPU Acceleration: Leverage GPU computing for high-performance simulation

Environments Module (envs)

Base Environment Classes

class embodichain.lab.gym.envs.BaseEnv

Bases: Env

Base environment for robot learning.

Parameters:

• cfg (EnvCfg) – The environment configuration.

• **kwargs – Additional keyword arguments.

The foundational environment class that provides the core functionality for all EmbodiChain RL environments. This class extends the Gymnasium Env interface with multi-environment support and robotic-specific features. Methods:

__init__(cfg, **kwargs)
add_camera_group_id(group_id)	Add a camera group ID for rendering.
add_detached_uids_for_reset(uids)	Add the UIDs of objects that are detached from automatic reset.
check_truncated(obs, info)	Check if the episode is truncated.
close()	Close the environment and release resources.
evaluate(**kwargs)	Evaluate whether the environment is currently in a success state by returning a dictionary with a "success" key or a failure state via a "fail" key
get_info(**kwargs)	Get info about the current environment state, include elapsed steps, success, fail, etc.
get_obs(**kwargs)	Get the observation from the robot agent and the environment.
get_reward(obs, action, info)	Get the reward for the current step.
get_sensor(name, **kwargs)	Get the sensor instance by name.
get_wrapper_attr(name)	Gets the attribute name from the environment.
has_wrapper_attr(name)	Checks if the attribute name exists in the environment.
is_task_success(**kwargs)	Determine if the task is successfully completed.
render()	Compute the render frames as specified by render_mode during the initialization of the environment.
reset([seed, options])	Reset the SimulationManager environment and return the observation and info.
set_wrapper_attr(name, value, *[, force])	Sets the attribute name on the environment with value, see Wrapper.set_wrapper_attr for more info.
step(action, **kwargs)	Step the environment with the given action.

Attributes:

control_frequency	Return the environment control frequency.
device	Return the device used by the environment.
flattened_observation_space	Flattened observation space for RL training.
has_sensors	Return whether the environment has sensors.
np_random	Returns the environment's internal _np_random that if not set will initialise with a random seed.
np_random_seed	Returns the environment's internal _np_random_seed that if not set will first initialise with a random int as seed.
num_envs	Return the number of environments simulated in parallel.
physics_dt	Return the duration of one physics simulation step.
physics_frequency	Return the physics simulation frequency.
step_dt	Return the duration of one environment control step.
unwrapped	Returns the base non-wrapped environment.

__init__(cfg, **kwargs)

add_camera_group_id(group_id)

Add a camera group ID for rendering.

Parameters:

group_id (int) – The camera group ID to be added.

Return type:

None

add_detached_uids_for_reset(uids)

Add the UIDs of objects that are detached from automatic reset.

Parameters:

uids (List[str]) – The list of UIDs to be detached from automatic reset.

Return type:

None

check_truncated(obs, info)

Check if the episode is truncated.

Parameters:

• obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation from the environment.

• info (TensorDict[str, Any]) – The info dictionary.

Return type:

Tensor

Returns:

A boolean tensor indicating truncation for each environment in the batch.

close()

Close the environment and release resources.

Return type:

None

property control_frequency: float

Return the environment control frequency.

Returns:

Environment control frequency in hertz.

property device: device

Return the device used by the environment.

evaluate(**kwargs)

Evaluate whether the environment is currently in a success state by returning a dictionary with a “success” key or a failure state via a “fail” key

This function may also return additional data that has been computed (e.g. is the robot grasping some object) that may be reused when generating observations and rewards.

By default if not overridden, this function returns an empty dictionary

Parameters:

**kwargs – Additional keyword arguments to be passed to the evaluate() function.

Return type:

Dict[str, Any]

Returns:

The evaluation dictionary.

property flattened_observation_space: Box

Flattened observation space for RL training.

Returns a Box space by computing total dimensions from nested dict observations. This is needed because RL algorithms (PPO, SAC, etc.) require flat vector inputs.

get_info(**kwargs)

Get info about the current environment state, include elapsed steps, success, fail, etc.

The returned info dictionary must contain at the success and fail status of the current step.

Parameters:

**kwargs – Additional keyword arguments to be passed to the get_info() function.

Return type:

TensorDict[str, Any]

Returns:

The info dictionary.

get_obs(**kwargs)

Get the observation from the robot agent and the environment.

The default observation are:

• robot: the robot proprioception.

• sensor (optional): the sensor readings.

• extra (optional): any extra information.

Parameters:

**kwargs – Additional keyword arguments to be passed to the _get_sensor_obs() functions.

Return type:

TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]

Returns:

The observation dictionary.

get_reward(obs, action, info)

Get the reward for the current step.

Each SimulationManager env must implement its own get_reward function to define the reward function for the task, If the env is considered for RL/IL training.

Parameters:

• obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation from the environment.

• action (Union[Tensor, TensorDict[str, Tensor]]) – The action applied to the robot agent.

• info (Dict[str, Any]) – The info dictionary.

Return type:

float

Returns:

The reward for the current step.

get_sensor(name, **kwargs)

Get the sensor instance by name.

Parameters:

• name (str) – The name of the sensor.

• kwargs – Additional keyword arguments.

Return type:

BaseSensor

Returns:

The sensor instance.

get_wrapper_attr(name)

Gets the attribute name from the environment.

Return type:

Any

property has_sensors: bool

Return whether the environment has sensors.

has_wrapper_attr(name)

Checks if the attribute name exists in the environment.

Return type:

bool

is_task_success(**kwargs)

Determine if the task is successfully completed. This is mainly used in the data generation process of the imitation learning.

Parameters:

**kwargs – Additional arguments for task-specific success criteria.

Returns:

A boolean tensor indicating success for each environment in the batch.

Return type:

torch.Tensor

property np_random: Generator

Returns the environment’s internal _np_random that if not set will initialise with a random seed.

Returns:

Instances of np.random.Generator

property np_random_seed: int

Returns the environment’s internal _np_random_seed that if not set will first initialise with a random int as seed.

If np_random_seed was set directly instead of through reset() or set_np_random_through_seed(), the seed will take the value -1.

Returns:

the seed of the current np_random or -1, if the seed of the rng is unknown

Return type:

int

property num_envs: int

Return the number of environments simulated in parallel.

property physics_dt: float

Return the duration of one physics simulation step.

Returns:

Physics simulation step duration in seconds.

property physics_frequency: float

Return the physics simulation frequency.

Returns:

Physics simulation frequency in hertz.

render()

Compute the render frames as specified by render_mode during the initialization of the environment.

The environment’s metadata render modes (env.metadata[“render_modes”]) should contain the possible ways to implement the render modes. In addition, list versions for most render modes is achieved through gymnasium.make which automatically applies a wrapper to collect rendered frames. :rtype: str | ndarray | tuple[ndarray, ndarray] | list[str | ndarray | tuple[ndarray, ndarray]] | None

Note

As the render_mode is known during __init__, the objects used to render the environment state should be initialised in __init__.

By convention, if the render_mode is:

• None (default): no render is computed.

• “human”: The environment is continuously rendered in the current display or terminal, usually for human consumption. This rendering should occur during step() and render() doesn’t need to be called. Returns None.

• “rgb_array”: Return a single frame representing the current state of the environment. A frame is a np.ndarray with shape (x, y, 3) representing RGB values for an x-by-y pixel image.

• “ansi”: Return a strings (str) or StringIO.StringIO containing a terminal-style text representation for each time step. The text can include newlines and ANSI escape sequences (e.g. for colors).

• “rgb_array_list” and “ansi_list”: List based version of render modes are possible (except Human) through the wrapper, gymnasium.wrappers.RenderCollection that is automatically applied during gymnasium.make(..., render_mode="rgb_array_list"). The frames collected are popped after render() is called or reset().

Note

Make sure that your class’s metadata "render_modes" key includes the list of supported modes.

Changed in version 0.25.0: The render function was changed to no longer accept parameters, rather these parameters should be specified in the environment initialised, i.e., gymnasium.make("CartPole-v1", render_mode="human")

reset(seed=None, options=None)

Reset the SimulationManager environment and return the observation and info.

Parameters:

• seed (int | None) – The seed for the random number generator. Defaults to None, in which case the seed is not set.

• options (dict | None) – Additional reset options. A trusted host may supply a zero-argument episode_prepare callback. It runs after reset events and before objective initialization and the final observation. The callback is not forwarded to observation or event functors.

Return type:

Tuple[TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]], Dict]

Returns:

A tuple containing the observations and infos.

set_wrapper_attr(name, value, *, force=True)

Sets the attribute name on the environment with value, see Wrapper.set_wrapper_attr for more info.

Return type:

bool

step(action, **kwargs)

Step the environment with the given action.

Parameters:

action (Union[Tensor, TensorDict[str, Tensor]]) – The action applied to the robot agent.

Return type:

Tuple[TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]], Tensor, Tensor, Tensor, Dict[str, Any]]

Returns:

A tuple contraining the observation, reward, terminated, truncated, and info dictionary.

property step_dt: float

Return the duration of one environment control step.

Returns:

Environment control step duration in seconds.

property unwrapped: Env[ObsType, ActType]

Returns the base non-wrapped environment.

Returns:

The base non-wrapped gymnasium.Env instance

Return type:

Env

class embodichain.lab.gym.envs.EnvCfg

Configuration for an Robot Learning Environment.

Configuration class for basic environment settings including simulation parameters and environment count. Methods:

copy(**kwargs)	Return a new object replacing specified fields with new values.
replace(**kwargs)	Return a new object replacing specified fields with new values.
to_dict()	Convert an object into dictionary recursively.
validate([prefix])	Check the validity of configclass object.

Attributes:

ignore_terminations	Whether to ignore terminations when deciding when to auto reset.
max_episode_steps	The maximum number of steps per episode.
num_envs	The number of sub environments (arena in dexsim context) to be simulated in parallel.
profiler	Optional profiler for reset/step wall-time breakdown.
seed	The task-environment seed.
sim_cfg	Simulation configuration for the environment.
sim_steps_per_control	Number of simulation steps per control (env) step.
target_control_frequency	Optional requested control frequency in hertz.

copy(**kwargs)

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2

Parameters:

• obj (object) – The object to replace.

• **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

ignore_terminations: bool

Whether to ignore terminations when deciding when to auto reset. Terminations can be caused by the task reaching a success or fail state as defined in a task’s evaluation function.

If set to False, meaning there is early stop in episode rollouts. If set to True, this would generally for situations where you may want to model a task as infinite horizon where a task stops only due to the timelimit.

max_episode_steps: int

The maximum number of steps per episode. If set to -1, there is no limit on the episode length, and the episode will only end when the task is successfully completed or failed.

num_envs: int

The number of sub environments (arena in dexsim context) to be simulated in parallel.

profiler: ProfilerCfg | None

Optional profiler for reset/step wall-time breakdown. None keeps the profiler disabled unless one is configured directly on sim_cfg. See EnvProfilerCfg for the available options.

replace(**kwargs)

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2

Parameters:

• obj (object) – The object to replace.

• **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

seed: int | None

The task-environment seed. Defaults to None, in which case the seed is not set.

Note

The seed is set before scene initialization and controls process RNGs and deterministic event-functor streams.

sim_cfg: SimulationManagerCfg

Simulation configuration for the environment.

sim_steps_per_control: int

Number of simulation steps per control (env) step.

For instance, if the simulation dt is 0.01s and the control dt is 0.1s, then the sim_steps_per_control is 10. This means that the control action is updated every 10 simulation steps.

target_control_frequency: float | None

Optional requested control frequency in hertz.

When set, the environment resolves this value to an integer sim_steps_per_control using the configured physics timestep and takes precedence over the directly configured step count. The requested frequency must be exactly representable; the physics timestep is never changed and the frequency is never silently approximated.

to_dict()

Convert an object into dictionary recursively.

Note

Ignores all names starting with “__” (i.e. built-in methods).

Parameters:

obj (object) – An instance of a class to convert.

Raises:

ValueError – When input argument is not an object.

Return type:

dict[str, Any]

Returns:

Converted dictionary mapping.

validate(prefix='')

Check the validity of configclass object.

This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.

Parameters:

• obj (object) – The object to check.

• prefix (str) – The prefix to add to the missing fields. Defaults to ‘’.

Return type:

list[str]

Returns:

A list of missing fields.

Raises:

TypeError – When the object is not a valid configuration object.

Embodied Environment Classes

class embodichain.lab.gym.envs.EmbodiedEnv

Bases: BaseEnv

Embodied AI environment that is used to simulate the Embodied AI tasks.

Core simulation components for Embodied AI environments. - sensor: The sensors used to perceive the environment, which could be attached to the agent or the environment. - robot: The robot which will be used to interact with the environment. - light: The lights in the environment, which could be used to illuminate the environment.

• indirect: the indirect light sources, such as ambient light, IBL, etc.

  The indirect light sources are used for global illumination which affects the entire scene.

• direct: The direct light sources, such as point light, spot light, etc.

  The direct light sources are used for local illumination which mainly affects the arena in the scene.

• background: Kinematic or Static rigid objects, such as obstacles or landmarks.

• rigid_object: Dynamic objects that can be interacted with.

• rigid_object_group: Groups of rigid objects that can be interacted with.

• deformable_object(TODO: supported in the future): Deformable volumes or surfaces (cloth) that can be interacted with.

• articulation: Articulated objects that can be manipulated, such as doors, drawers, etc.

• event manager: The event manager is used to manage the events in the environment, such as randomization,

  perturbation, etc.

• observation manager: The observation manager is used to manage the observations in the environment,

  such as depth, segmentation, etc.

• affordance_datas: The affordance data that can be used to store the intermediate results or information

An advanced environment class that provides additional features for embodied AI research, including sophisticated observation management, event handling, and multi-modal sensor integration. Methods:

__init__(cfg, *[, task_program_adapter_factory])
add_camera_group_id(group_id)	Add a camera group ID for rendering.
add_detached_uids_for_reset(uids)	Add the UIDs of objects that are detached from automatic reset.
check_truncated(obs, info)	Check if the episode is truncated.
close(*[, exit_process])	Abort pending data, finalize committed writes, and release resources.
compile_task_program(program)	Compile a configured Task Program through the explicit adapter.
compute_task_state(**kwargs)	Compute task-specific state: success, fail, and metrics.
create_demo_action_list(*args, **kwargs)	Create a demonstration action list for the environment.
create_demo_segments(*args[, task_program, ...])	Create the semantic segments that make up one task episode.
create_task_program_bridge(program, *[, ...])	Create the Gym demo bridge through the explicit adapter.
evaluate(**kwargs)	Evaluate the environment state.
get_demo_episode_metadata(env_id)	Return segment-aware metadata for one buffered episode.
get_executed_qpos_history(env_id[, length])	Compatibility alias for expert controller qpos history.
get_expert_controller_qpos_history(env_id[, ...])	Return expert controller qpos commands for one episode.
get_info(**kwargs)	Get environment info dictionary.
get_obs(**kwargs)	Get the observation from the robot agent and the environment.
get_raw_action_history(env_id[, length])	Return raw actions recorded for one environment episode.
get_reward(obs, action, info)	Get the reward for the current step.
get_sensor(name, **kwargs)	Get the sensor instance by name.
get_wrapper_attr(name)	Gets the attribute name from the environment.
has_wrapper_attr(name)	Checks if the attribute name exists in the environment.
is_task_success(**kwargs)	Return completed Task Program acceptance or legacy task success.
preview_sensor_data(name[, data_type, ...])	Preview the sensor data by matplotlib
render()	Compute the render frames as specified by render_mode during the initialization of the environment.
reset([seed, options])	Reset the SimulationManager environment and return the observation and info.
save_trajectory(path[, env_ids, purpose])	Save a causally aligned trajectory to a .pt file.
set_rollout_buffer(rollout_buffer)	Set the rollout buffer for episode data collection.
set_wrapper_attr(name, value, *[, force])	Sets the attribute name on the environment with value, see Wrapper.set_wrapper_attr for more info.
step(action, **kwargs)	Step the environment with the given action.

Attributes:

control_frequency	Return the environment control frequency.
device	Return the device used by the environment.
flattened_observation_space	Flattened observation space for RL training.
has_sensors	Return whether the environment has sensors.
last_action_trace	Return the most recent ActionManager request/command trace.
np_random	Returns the environment's internal _np_random that if not set will initialise with a random seed.
np_random_seed	Returns the environment's internal _np_random_seed that if not set will first initialise with a random int as seed.
num_envs	Return the number of environments simulated in parallel.
physics_dt	Return the duration of one physics simulation step.
physics_frequency	Return the physics simulation frequency.
rollout_action_semantics	Return the explicit meaning of rollout_buffer['actions'].
step_dt	Return the duration of one environment control step.
task_program_adapter	Return the adapter injected after the environment built its scene.
task_program_expansion_records	Return records owned by the active generated Task Program episode.
unwrapped	Returns the base non-wrapped environment.

__init__(cfg, *, task_program_adapter_factory=None, **kwargs)

add_camera_group_id(group_id)

Add a camera group ID for rendering.

Parameters:

group_id (int) – The camera group ID to be added.

Return type:

None

add_detached_uids_for_reset(uids)

Add the UIDs of objects that are detached from automatic reset.

Parameters:

uids (List[str]) – The list of UIDs to be detached from automatic reset.

Return type:

None

check_truncated(obs, info)

Check if the episode is truncated.

Parameters:

• obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation from the environment.

• info (TensorDict[str, Any]) – The info dictionary.

Return type:

Tensor

Returns:

A boolean tensor indicating truncation for each environment in the batch.

close(*, exit_process=None)

Abort pending data, finalize committed writes, and release resources.

Closing is idempotent and is never an implicit episode commit. A demo episode enters the dataset only through reset(save_data=True); partial data left by failure, cancellation, or interpreter shutdown is discarded before recorder finalization.

Parameters:

exit_process (bool | None) – Forwarded to SimulationManager.destroy() after successful cleanup. Error paths always disable process exit so the durability exception can propagate.

Raises:

RuntimeError – If one or more recorders fail their durability barrier.

Return type:

None

compile_task_program(program)

Compile a configured Task Program through the explicit adapter.

Parameters:

program (TaskProgramCfg) – Strict Task Program configuration attached to cfg.

Return type:

CompiledTaskProgram

Returns:

Provider-free compiled program ready for runtime assembly.

compute_task_state(**kwargs)

Compute task-specific state: success, fail, and metrics.

Override this method in subclass to define task-specific logic for RL tasks.

Returns:

Success and failure masks of shape (num_envs,), followed by the metric tensor mapping.

Return type:

tuple[torch.Tensor, torch.Tensor, dict[str, Any]]

property control_frequency: float

Return the environment control frequency.

Returns:

Environment control frequency in hertz.

create_demo_action_list(*args, **kwargs)

Create a demonstration action list for the environment.

This function should be implemented in subclasses to generate a sequence of actions that demonstrate a specific task or behavior within the environment.

Returns:

A list of actions if a demonstration is available, otherwise None.

Return type:

Sequence[EnvAction] | None

Note

Subclass outputs are automatically post-processed by the base class: action last-dimension must match single_action_space. If larger, actions are sliced by active_joint_ids; if smaller, ValueError is raised.

create_demo_segments(*args, task_program=None, expansion_profile=None, expansion_candidate_index=0, expansion_recipe_indices=None, **kwargs)

Create the semantic segments that make up one task episode.

An episode-level task_program takes precedence over the static configuration. This lets trusted callers supply a model-produced, already compiled program without mutating cfg. Otherwise, a configured program is compiled through the injected adapter. With no selected program, the legacy action-list path remains unchanged.

Parameters:

• *args – Positional arguments forwarded to the legacy planner.

• task_program (TaskProgramCfg | CompiledTaskProgram | None) – Optional episode-level program config or provider-free compiled program.

• expansion_profile (CombinedExpansionProfile | None) – Optional source-neutral profile applied to matching grounded Atomic calls in this episode.

• expansion_candidate_index (int) – FIFO candidate ordinal selected from each generated reference.

• **kwargs – Keyword arguments forwarded to the legacy planner.

Return type:

Optional[Iterable[DemoSegment]]

Returns:

Segment sequence, or None when planning fails.

create_task_program_bridge(program, *, expansion_profile=None, expansion_candidate_index=0, expansion_recipe_indices=None)

Create the Gym demo bridge through the explicit adapter.

Parameters:

• program (CompiledTaskProgram) – Compiled provider-free Task Program.

• expansion_profile (CombinedExpansionProfile | None) – Optional episode-scoped Expansion Profile.

• expansion_candidate_index (int) – FIFO candidate ordinal selected from each matching Atomic plan.

Return type:

TaskProgramDemoBridge

Returns:

Atomic demo bridge whose segments are consumed lazily.

property device: device

Return the device used by the environment.

evaluate(**kwargs)

Evaluate the environment state.

Return type:

Dict[str, Any]

Returns:

Evaluation dictionary with success and metrics

property flattened_observation_space: Box

Flattened observation space for RL training.

Returns a Box space by computing total dimensions from nested dict observations. This is needed because RL algorithms (PPO, SAC, etc.) require flat vector inputs.

get_demo_episode_metadata(env_id)

Return segment-aware metadata for one buffered episode.

Legacy collection paths that do not use the common executor are represented as one segment spanning every valid frame.

Parameters:

env_id (int) – Parallel environment row.

Return type:

dict[str, Any]

Returns:

A JSON-compatible metadata dictionary.

get_executed_qpos_history(env_id, length=None)

Compatibility alias for expert controller qpos history.

Return type:

Tensor

get_expert_controller_qpos_history(env_id, length=None)

Return expert controller qpos commands for one episode.

Return type:

Tensor

get_info(**kwargs)

Get environment info dictionary.

Calls compute_task_state() to get task-specific success/fail/metrics when available. Subclasses should override compute_task_state() for RL tasks.

Return type:

Dict[str, Any]

Returns:

Info dictionary with success, fail, elapsed_steps, metrics

get_obs(**kwargs)

Get the observation from the robot agent and the environment.

The default observation are:

• robot: the robot proprioception.

• sensor (optional): the sensor readings.

• extra (optional): any extra information.

Parameters:

**kwargs – Additional keyword arguments to be passed to the _get_sensor_obs() functions.

Return type:

TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]

Returns:

The observation dictionary.

get_raw_action_history(env_id, length=None)

Return raw actions recorded for one environment episode.

Return type:

Any

get_reward(obs, action, info)

Get the reward for the current step.

Each SimulationManager env must implement its own get_reward function to define the reward function for the task, If the env is considered for RL/IL training.

Parameters:

• obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation from the environment.

• action (Union[Tensor, TensorDict[str, Tensor]]) – The action applied to the robot agent.

• info (Dict[str, Any]) – The info dictionary.

Return type:

float

Returns:

The reward for the current step.

get_sensor(name, **kwargs)

Get the sensor instance by name.

Parameters:

• name (str) – The name of the sensor.

• kwargs – Additional keyword arguments.

Return type:

BaseSensor

Returns:

The sensor instance.

get_wrapper_attr(name)

Gets the attribute name from the environment.

Return type:

Any

property has_sensors: bool

Return whether the environment has sensors.

has_wrapper_attr(name)

Checks if the attribute name exists in the environment.

Return type:

bool

is_task_success(**kwargs)

Return completed Task Program acceptance or legacy task success.

Task Program success is published only after its bridge has consumed every segment lifecycle, including post-policies and validators.

Parameters:

**kwargs (Any) – Compatibility keywords forwarded for tasks without a Task Program.

Return type:

Tensor

Returns:

Per-environment task-success mask.

property last_action_trace: ActionTrace | None

Return the most recent ActionManager request/command trace.

property np_random: Generator

Returns the environment’s internal _np_random that if not set will initialise with a random seed.

Returns:

Instances of np.random.Generator

property np_random_seed: int

Returns the environment’s internal _np_random_seed that if not set will first initialise with a random int as seed.

If np_random_seed was set directly instead of through reset() or set_np_random_through_seed(), the seed will take the value -1.

Returns:

the seed of the current np_random or -1, if the seed of the rng is unknown

Return type:

int

property num_envs: int

Return the number of environments simulated in parallel.

property physics_dt: float

Return the duration of one physics simulation step.

Returns:

Physics simulation step duration in seconds.

property physics_frequency: float

Return the physics simulation frequency.

Returns:

Physics simulation frequency in hertz.

preview_sensor_data(name, data_type='color', env_ids=0, method='cv2', save=False)

Preview the sensor data by matplotlib

Note

Currently only support RGB image preview.

Parameters:

• name (str) – name of the sensor to preview.

• data_type (str) – type of the sensor data to preview.

• env_ids (int) – index of the arena to preview. Defaults to 0.

• method (str) – method to preview the sensor data. Currently support “plt” and “cv2”. Defaults to “cv2”.

• save (bool) – whether to save the preview image. Defaults to False.

Return type:

None

render()

Compute the render frames as specified by render_mode during the initialization of the environment.

The environment’s metadata render modes (env.metadata[“render_modes”]) should contain the possible ways to implement the render modes. In addition, list versions for most render modes is achieved through gymnasium.make which automatically applies a wrapper to collect rendered frames. :rtype: str | ndarray | tuple[ndarray, ndarray] | list[str | ndarray | tuple[ndarray, ndarray]] | None

Note

As the render_mode is known during __init__, the objects used to render the environment state should be initialised in __init__.

By convention, if the render_mode is:

• None (default): no render is computed.

• “human”: The environment is continuously rendered in the current display or terminal, usually for human consumption. This rendering should occur during step() and render() doesn’t need to be called. Returns None.

• “rgb_array”: Return a single frame representing the current state of the environment. A frame is a np.ndarray with shape (x, y, 3) representing RGB values for an x-by-y pixel image.

• “ansi”: Return a strings (str) or StringIO.StringIO containing a terminal-style text representation for each time step. The text can include newlines and ANSI escape sequences (e.g. for colors).

• “rgb_array_list” and “ansi_list”: List based version of render modes are possible (except Human) through the wrapper, gymnasium.wrappers.RenderCollection that is automatically applied during gymnasium.make(..., render_mode="rgb_array_list"). The frames collected are popped after render() is called or reset().

Note

Make sure that your class’s metadata "render_modes" key includes the list of supported modes.

Changed in version 0.25.0: The render function was changed to no longer accept parameters, rather these parameters should be specified in the environment initialised, i.e., gymnasium.make("CartPole-v1", render_mode="human")

reset(seed=None, options=None)

Reset the SimulationManager environment and return the observation and info.

Parameters:

• seed (int | None) – The seed for the random number generator. Defaults to None, in which case the seed is not set.

• options (dict | None) – Additional reset options. A trusted host may supply a zero-argument episode_prepare callback. It runs after reset events and before objective initialization and the final observation. The callback is not forwarded to observation or event functors.

Return type:

Tuple[TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]], Dict]

Returns:

A tuple containing the observations and infos.

property rollout_action_semantics: str | None

Return the explicit meaning of rollout_buffer['actions'].

Values are policy_request for RL/raw-policy buffers and expert_controller_command for expert/demo buffers.

save_trajectory(path, env_ids=None, *, purpose='recording')

Save a causally aligned trajectory to a .pt file.

states[t] is the state immediately before actions[t] is applied, matching the frame alignment used by expert/LeRobot trajectories.

Parameters:

• path (str) – Destination .pt file path.

• env_ids (Sequence[int] | None) – Env indices to save (default: all). Each saved env’s actual recorded length is stored in meta["lengths"].

• purpose (str) – Normal recording or explicitly labeled diagnostic export.

Raises:

RuntimeError – If trajectory recording was never enabled.

Return type:

str

set_rollout_buffer(rollout_buffer)

Set the rollout buffer for episode data collection.

This function can be used to set the rollout buffer from outside of the environment, such as a shared rollout buffer initialized in model training process and passed to the environment for data collection.

Parameters:

rollout_buffer (TensorDict) – The rollout buffer to be set. RL rollouts use a uniform [num_envs, time + 1] layout so all fields share the same batch shape; the last slot of transition-only fields is reserved as padding. Expert buffers keep the legacy [num_envs, time] batch layout.

Return type:

None

set_wrapper_attr(name, value, *, force=True)

Sets the attribute name on the environment with value, see Wrapper.set_wrapper_attr for more info.

Return type:

bool

step(action, **kwargs)

Step the environment with the given action.

Parameters:

action (Union[Tensor, TensorDict[str, Tensor]]) – The action applied to the robot agent.

Return type:

Tuple[TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]], Tensor, Tensor, Tensor, Dict[str, Any]]

Returns:

A tuple contraining the observation, reward, terminated, truncated, and info dictionary.

property step_dt: float

Return the duration of one environment control step.

Returns:

Environment control step duration in seconds.

property task_program_adapter: TaskProgramEnvironmentAdapter

Return the adapter injected after the environment built its scene.

A registered environment normally receives an TaskProgramAdapterFactory through its EnvSpec. Advanced integrations may still override this property.

property task_program_expansion_records: tuple[TaskProgramExpansionRecord, ...]

Return records owned by the active generated Task Program episode.

property unwrapped: Env[ObsType, ActType]

Returns the base non-wrapped environment.

Returns:

The base non-wrapped gymnasium.Env instance

Return type:

Env

class embodichain.lab.gym.envs.EmbodiedEnvCfg

Configuration for Embodied AI environments.

EmbodiedEnvCfg extends EnvCfg with high-level scene, robot, sensor, object and manager declarations used to build modular embodied environments. The configuration is intended to be declarative: the environment and its managers (events, observations, rewards, dataset) are assembled from the provided config fields with minimal additional code.

Typical usage: declare robots, sensors, lights, rigid objects/articulations, and manager configurations. Additional task-specific parameters can be supplied via the extensions dict and will be bound to the environment instance as attributes during initialization.

Key fields - robot: RobotCfg | RobotPresetCfg (required) — one portable robot

definition or replace-only complete alternatives selected by the active physics backend.

• control_parts: Optional[List[str]] — named robot parts to control. If

  None, all controllable joints are used.

• active_joint_ids: List[int] — explicit joint indices to use for

  control (alternative to control_parts).

• sensor: List[SensorCfg] — independent sensors. Mounted cameras

  declared by robot.sensors are included in the effective sensor suite.

• light: EnvLightCfg — lighting configuration (direct lights now,

  indirect/IBL planned for future releases).

• background, rigid_object, rigid_object_group, articulation:

  scene object lists for static/kinematic props, dynamic objects, grouped object pools, and articulated mechanisms respectively.

• events: Optional manager config — event functors for startup/reset/

  periodic randomization and scripted behaviors.

• observations, rewards, dataset: Optional manager configs to

  compose observation transforms, reward functors, and dataset/recorder settings (auto-saving on episode completion).

• extensions: Optional[Dict[str, Any]] — arbitrary task-specific key/value

  pairs (e.g. success_threshold, control_frequency) that are automatically set on the config and bound to the environment instance.

• filter_visual_rand / filter_dataset_saving: booleans to disable

  visual randomization or dataset saving for debugging purposes.

• init_rollout_buffer: bool — when true (or when a dataset manager is

  present and dataset saving is enabled) the environment will initialize a rollout buffer matching the observation/action spaces for episode recording.

See EmbodiedEnv for usage patterns and the project documentation for full examples showing how to declare environments from these configs.

Configuration class for embodied environments with extended settings for lighting, observation management, and advanced simulation features. Classes:

EnvLightCfg

EnvLightCfg(direct: 'List[LightCfg]' = <factory>, indirect: 'dict[str, Any] | None' = <factory>)

Attributes:

actions	Action manager settings.
active_joint_ids	List of active joint IDs for control.
control_parts	List of robot parts to control.
dataset	Dataset settings.
enable_sensor	Whether configured sensors are instantiated and sampled.
events	Event settings.
expert_trajectory	Source-neutral expert trajectory control and recording settings.
extensions	Extension parameters for task-specific configurations.
filter_dataset_saving	Whether to filter out dataset saving
filter_visual_rand	Whether to filter out visual randomization
ground_plane_visible	Native rendering visibility of the default ground plane.
ignore_terminations	Whether to ignore terminations when deciding when to auto reset.
init_rollout_buffer	Whether to initialize the rollout buffer in the environment.
max_episode_steps	The maximum number of steps per episode.
num_envs	The number of sub environments (arena in dexsim context) to be simulated in parallel.
objective	Optional measured physical objective, independent of task success and persistence.
observations	Observation settings.
profiler	Optional profiler for reset/step wall-time breakdown.
record_trajectory	Whether to record per-object states and pre-process actions.
rewards	Reward settings.
rollout_action_semantics	Meaning of episode rollout actions, independently of the expert schema.
seed	The task-environment seed.
sensor	Independent sensors, combined with the robot's mounted sensor suite.
sim_cfg	Simulation configuration for the environment.
sim_steps_per_control	Number of simulation steps per control (env) step.
target_control_frequency	Optional requested control frequency in hertz.
task_instruction	Overall task language for handwritten/direct-control demonstrations.
task_program	Optional declarative Task Program used to generate demo segments.
trajectory_auto_save	If True (and record_trajectory is True), auto-save each env's trajectory to trajectory_save_dir before an explicitly committed episode reset.
trajectory_save_dir	Directory for auto-saved trajectories.
trajectory_uids	Optional allow-list of non-robot object uids to record.

Methods:

copy(**kwargs)	Return a new object replacing specified fields with new values.
replace(**kwargs)	Return a new object replacing specified fields with new values.
to_dict()	Convert an object into dictionary recursively.
validate([prefix])	Check the validity of configclass object.

class EnvLightCfg

EnvLightCfg(direct: ‘List[LightCfg]’ = <factory>, indirect: ‘dict[str, Any] | None’ = <factory>)

Methods:

copy(**kwargs)	Return a new object replacing specified fields with new values.
replace(**kwargs)	Return a new object replacing specified fields with new values.
to_dict()	Convert an object into dictionary recursively.
validate([prefix])	Check the validity of configclass object.

copy(**kwargs)

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2

Parameters:

• obj (object) – The object to replace.

• **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

replace(**kwargs)

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2

Parameters:

• obj (object) – The object to replace.

• **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

to_dict()

Convert an object into dictionary recursively.

Note

Ignores all names starting with “__” (i.e. built-in methods).

Parameters:

obj (object) – An instance of a class to convert.

Raises:

ValueError – When input argument is not an object.

Return type:

dict[str, Any]

Returns:

Converted dictionary mapping.

validate(prefix='')

Check the validity of configclass object.

This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.

Parameters:

• obj (object) – The object to check.

• prefix (str) – The prefix to add to the missing fields. Defaults to ‘’.

Return type:

list[str]

Returns:

A list of missing fields.

Raises:

TypeError – When the object is not a valid configuration object.

actions: Union[object, None]

Action manager settings. Defaults to None, in which case no action preprocessing is applied.

When configured, the ActionManager preprocesses raw policy actions (e.g., delta_qpos, eef_pose) into robot control format.

Please refer to the embodichain.lab.gym.envs.managers.ActionManager class for more details.

active_joint_ids: List[int]

List of active joint IDs for control. User also can directly specify the active joint IDs instead of control parts. This is useful when the control parts are not well defined or we want to have more fine-grained control.

control_parts: list[str] | None

List of robot parts to control. If None, all controllable joints will be used. This is useful when we want to control only a subset of the robot joints for certain tasks or demonstrations.

copy(**kwargs)

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2

Parameters:

• obj (object) – The object to replace.

• **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

dataset: Union[object, None]

Dataset settings. Defaults to None, in which case no dataset collection is performed.

Please refer to the embodichain.lab.gym.managers.DatasetManager class for more details.

enable_sensor: bool

Whether configured sensors are instantiated and sampled.

Independent declarations and robot-mounted defaults remain available when a deployment only needs proprioception. Set this to False to skip all sensor creation, including camera-group allocation and image fetching.

events: Union[object, None]

Event settings. Defaults to None, in which case no events are applied through the event manager.

Please refer to the embodichain.lab.gym.managers.EventManager class for more details.

expert_trajectory: ExpertTrajectoryCfg

Source-neutral expert trajectory control and recording settings.

extensions: Union[Dict[str, Any], None]

Extension parameters for task-specific configurations.

This field can be used to pass additional parameters that are specific to certain environments or tasks without modifying the base configuration class. For example: - success_threshold: Task-specific success distance threshold - vr_joint_mapping: VR joint mapping for teleoperation - control_frequency: Control frequency for VR teleoperation

Note: Action configuration (e.g., delta_qpos, scale) should use the actions field and ActionManager, not extensions.

filter_dataset_saving: bool

Whether to filter out dataset saving

This is useful when we want to disable dataset saving for debug motion and physics issues. If no dataset manager is configured, this flag will have no effect.

filter_visual_rand: bool

Whether to filter out visual randomization

This is useful when we want to disable visual randomization for debug motion and physics issues.

ground_plane_visible: bool | None

Native rendering visibility of the default ground plane.

None hides the ground when an active HDR indirect-lighting randomizer is configured, and shows it otherwise. Explicit booleans override that default. Visibility affects native windows and cameras without changing ground collisions or the Viser grid.

ignore_terminations: bool

Whether to ignore terminations when deciding when to auto reset. Terminations can be caused by the task reaching a success or fail state as defined in a task’s evaluation function.

If set to False, meaning there is early stop in episode rollouts. If set to True, this would generally for situations where you may want to model a task as infinite horizon where a task stops only due to the timelimit.

init_rollout_buffer: bool

Whether to initialize the rollout buffer in the environment.

If filter_dataset_saving is False and a dataset manager is configured, the rollout buffer will be initialized by default

max_episode_steps: int

The maximum number of steps per episode. If set to -1, there is no limit on the episode length, and the episode will only end when the task is successfully completed or failed.

num_envs: int

The number of sub environments (arena in dexsim context) to be simulated in parallel.

objective: OrderedPlacementObjectiveCfg | None

Optional measured physical objective, independent of task success and persistence.

observations: Union[object, None]

Observation settings. Defaults to None, in which case no additional observations are applied through the observation manager.

Please refer to the embodichain.lab.gym.managers.ObservationManager class for more details.

profiler: EnvProfilerCfg | None

Optional profiler for reset/step wall-time breakdown. None keeps the profiler disabled unless one is configured directly on sim_cfg. See EnvProfilerCfg for the available options.

record_trajectory: bool

Whether to record per-object states and pre-process actions.

Each saved row is a causal (state_t, action_t) pair, matching expert trajectory frame alignment. Uses a per-env step counter so async parallel environments are supported.

replace(**kwargs)

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2

Parameters:

• obj (object) – The object to replace.

• **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

rewards: Union[object, None]

Reward settings. Defaults to None, in which case no reward computation is performed through the reward manager.

Please refer to the embodichain.lab.gym.managers.RewardManager class for more details.

rollout_action_semantics: Literal['expert_controller_command', 'policy_request']

Meaning of episode rollout actions, independently of the expert schema.

policy_request stores the flat ActionManager input and requires an ActionManager. It also enables raw-policy recording without a DatasetManager. The default preserves active-joint expert-controller recording.

seed: int | None

The task-environment seed. Defaults to None, in which case the seed is not set.

Note

The seed is set before scene initialization and controls process RNGs and deterministic event-functor streams.

sensor: List[SensorCfg]

Independent sensors, combined with the robot’s mounted sensor suite.

sim_cfg: SimulationManagerCfg

Simulation configuration for the environment.

sim_steps_per_control: int

Number of simulation steps per control (env) step.

For instance, if the simulation dt is 0.01s and the control dt is 0.1s, then the sim_steps_per_control is 10. This means that the control action is updated every 10 simulation steps.

target_control_frequency: float | None

Optional requested control frequency in hertz.

When set, the environment resolves this value to an integer sim_steps_per_control using the configured physics timestep and takes precedence over the directly configured step count. The requested frequency must be exactly representable; the physics timestep is never changed and the frequency is never silently approximated.

task_instruction: str | None

Overall task language for handwritten/direct-control demonstrations.

Task Programs own their instruction in the program configuration. Explicit episode instructions take precedence; recorder instructions are legacy fallbacks and are not required to enable task language.

task_program: TaskProgramCfg | None

Optional declarative Task Program used to generate demo segments.

The program remains inert until EmbodiedEnv.create_demo_segments() requests an explicit environment compiler and bridge through the dedicated hooks. No live provider, planner, or callable is stored in this config.

to_dict()

Convert an object into dictionary recursively.

Note

Ignores all names starting with “__” (i.e. built-in methods).

Parameters:

obj (object) – An instance of a class to convert.

Raises:

ValueError – When input argument is not an object.

Return type:

dict[str, Any]

Returns:

Converted dictionary mapping.

trajectory_auto_save: bool

If True (and record_trajectory is True), auto-save each env’s trajectory to trajectory_save_dir before an explicitly committed episode reset. Closing discards trajectories that were never committed.

trajectory_save_dir: str | None

Directory for auto-saved trajectories. Defaults to <EMBODICHAIN_DEFAULT_DATA_ROOT>/trajectories/{run_id}/.

trajectory_uids: list[str] | None

Optional allow-list of non-robot object uids to record. If None, all rigid objects and articulations are recorded. The robot is always recorded.

validate(prefix='')

Check the validity of configclass object.

This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.

Parameters:

• obj (object) – The object to check.

• prefix (str) – The prefix to add to the missing fields. Defaults to ‘’.

Return type:

list[str]

Returns:

A list of missing fields.

Raises:

TypeError – When the object is not a valid configuration object.

Utilities Module (utils)

Registration System

Gym environment registration and task-package discovery utilities.

class embodichain.lab.gym.utils.registration.EnvSpec

Bases: object

Specification class for environment registration, containing environment metadata and creation parameters. Methods:

__init__(uid, cls[, max_episode_steps, ...])

A specification for a Embodied environment.

Attributes:

gym_spec

Return a gym EnvSpec for this env

__init__(uid, cls, max_episode_steps=None, default_kwargs=None, task_program_registration=None, task_program_adapter_factory=None, supports_rl=False)

A specification for a Embodied environment.

property gym_spec

Return a gym EnvSpec for this env

embodichain.lab.gym.utils.registration.register(name, cls, max_episode_steps=None, default_kwargs=None, task_program_registration=None, task_program_adapter_factory=None, supports_rl=False)

Register a Embodied environment.

Register a new environment class with the EmbodiChain environment registry.

Parameters:

• name – Unique identifier for the environment

• cls – Environment class (must inherit from BaseEnv or BaseEnv)

• default_kwargs – Default keyword arguments for environment creation

embodichain.lab.gym.utils.registration.register_env(uid, max_episode_steps=None, override=False, *, supports_rl=False, task_program_registration=None, task_program_adapter_factory=None, **kwargs)

A decorator to register Embodied environments.

Parameters:

• uid (str) – unique id of the environment.

• max_episode_steps (int) – maximum number of steps in an episode.

• override (bool) – whether to override the environment if it is already registered.

• supports_rl (bool) – Whether the environment has a supported RL training path.

Notes

• max_episode_steps is processed differently from other keyword arguments in gym. gym.make wraps the env with gym.wrappers.TimeLimit to limit the maximum number of steps.

• gym.EnvSpec uses kwargs instead of **kwargs!

Decorator function for registering environment classes. This is the recommended way to register environments.

Parameters:

• uid – Unique identifier for the environment

• override – Whether to override existing environment with same ID

• kwargs – Additional registration parameters

Example:

@register_env("MyEnv-v1")
class MyCustomEnv(BaseEnv):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

embodichain.lab.gym.utils.registration.make(env_id, **kwargs)

Instantiate a Embodied environment.

Parameters:

• env_id (str) – Environment ID.

• as_gym (bool, optional) – Add TimeLimit wrapper as gym.

• **kwargs – Keyword arguments to pass to the environment.

Create an environment instance from a registered environment ID.

Parameters:

• env_id – Registered environment identifier

• kwargs – Additional keyword arguments for environment creation

Returns:

Environment instance

class embodichain.lab.gym.utils.registration.TimeLimitWrapper

Bases: Wrapper

like the standard gymnasium timelimit wrapper but fixes truncated variable to be a batched array

Gymnasium wrapper that adds episode time limits to environments. Methods:

__init__(env, max_episode_steps)	Wraps an environment to allow a modular transformation of the step() and reset() methods.
step(action)	Uses the step() of the env that can be overwritten to change the returned data.

__init__(env, max_episode_steps)

Wraps an environment to allow a modular transformation of the step() and reset() methods.

Parameters:

env (Env) – The environment to wrap

step(action)

Uses the step() of the env that can be overwritten to change the returned data.

Gymnasium Utilities

Helper functions and utilities for Gymnasium environment integration.

Miscellaneous Utilities

Miscellaneous utility functions for environment development and debugging.