embodichain.lab.gym.envs

embodichain.lab.gym.envs#

Environment framework: BaseEnv / EmbodiedEnv class hierarchy, task registration, manager wiring, and the step/reset lifecycle.

Overview#

The Gymnasium-compatible environment framework for embodied manipulation tasks. BaseEnv extends gym.Env with multi-environment (vectorized) support and owns the SimulationManager, robot, sensors, and action/observation spaces. EmbodiedEnv builds on BaseEnv and is the modular base class for concrete tasks: it wires in the event, observation, reward, action, and dataset managers via the functor/FunctorCfg pattern. Tasks are registered with register_env() and instantiated through make().

Submodules

demo

Segment-aware expert demonstration protocol and executor.

differentiable_env

Newton-backed kinematic environment for analytic policy gradient.

expert_trajectory

Expert action encoding at the Gym and dataset boundary.

task_program

Gym lifecycle bridge for compiled Task Programs.

managers

Managers that orchestrate collections of functors (observation, reward, event, action, dataset) running at specific points in the environment step loop.

types

Shared action types for the Gym environment boundary.

wrapper

Environment Classes#

class embodichain.lab.gym.envs.BaseEnv[source]#

Bases: Env

Base environment for robot learning.

Parameters:

cfg (EnvCfg) – The environment configuration.
**kwargs – Additional keyword arguments.

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_generator`(name)	Return a named component RNG rewound by explicit reset seeds.
`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)[source]#

add_camera_group_id(group_id)[source]#

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

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

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

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

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

Return a named component RNG rewound by explicit reset seeds.

Each name has its own stream on the environment device. Ordinary selective resets consume the existing stream without rewinding it.

Parameters:: name (str) – Stable name identifying the owning component.
Return type:: Generator
Returns:: Persistent Torch generator for this component.

get_info(**kwargs)[source]#

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

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

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

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

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

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

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[source]#

Configuration for an Robot Learning Environment.

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.

class embodichain.lab.gym.envs.EmbodiedEnv[source]#

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

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_generator`(name)	Return a named component RNG rewound by explicit reset seeds.
`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)[source]#

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

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

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

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

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

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

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

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

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

Compatibility alias for expert controller qpos history.

Return type:: Tensor

get_expert_controller_qpos_history(env_id, length=None)[source]#

Return expert controller qpos commands for one episode.

Return type:: Tensor

get_generator(name)#

Return a named component RNG rewound by explicit reset seeds.

Each name has its own stream on the environment device. Ordinary selective resets consume the existing stream without rewinding it.

Parameters:: name (str) – Stable name identifying the owning component.
Return type:: Generator
Returns:: Persistent Torch generator for this component.

get_info(**kwargs)[source]#

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

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

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

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

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

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[source]#

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.

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.
`visual_object`	Render-only assets with no collision or rigid-body state.

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[source]#

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.

visual_object: list[VisualObjectCfg]#: Render-only assets with no collision or rigid-body state.

Differentiable Environment#

DifferentiableEnv keeps the standard environment lifecycle while bridging task-defined Newton kinematics into PyTorch autograd for analytic policy-gradient tasks. Subclasses provide action, kinematics, and output kernels; the base class owns tape-aware stepping and deferred resets without advancing the Newton solver.

class embodichain.lab.gym.envs.differentiable_env.DifferentiableEnv[source]#

Bases: EmbodiedEnv

Newton-only environment with an APG-ready kinematic step().

Subclasses implement _apply_action_kernel(), _make_kinematic_step_fn(), and _read_outputs(). The action, kinematics, observation, and reward kernels execute inside one Warp tape. No Newton solver or collision step is invoked by this environment.

Methods:

`__init__`(cfg, args, *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`(*[, 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_generator`(name)	Return a named component RNG rewound by explicit reset seeds.
`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)	Advance one differentiable control step.

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, *args, **kwargs)[source]#

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_generator(name)#

Return a named component RNG rewound by explicit reset seeds.

Each name has its own stream on the environment device. Ordinary selective resets consume the existing stream without rewinding it.

Parameters:: name (str) – Stable name identifying the owning component.
Return type:: Generator
Returns:: Persistent Torch generator for this component.

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

Advance one differentiable control step.

Terminal environments are auto-reset only when the call cannot retain a Warp tape for backward. A grad-tracked step returns terminal observations unchanged and records deferred_reset_ids in info; callers must run backward before resetting those environments.

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

Expert Trajectories#

Expert trajectory configuration is source-neutral: handwritten generators, motion generation, and Task Program execution use the same environment-owned joint command mode. Position-only control remains the default. The optional position-velocity mode records a stable flat [qpos, qvel] action layout without changing the policy-facing Gym action space.

`EXPERT_TRAJECTORY_SCHEMA_VERSION`	int([x]) -> integer int(x, base=10) -> integer
`ExpertActionSpec`	Canonical flat action layout used by expert trajectory persistence.
`ExpertTrajectoryCfg`	Environment-owned expert trajectory execution and recording settings.
`JointCommandMode`	alias of `Literal`['position', 'position_velocity']
`build_expert_action_spec`	Build the immutable stored-action layout for one environment.
`encode_expert_action`	Flatten effective qpos/qvel targets into the expert dataset layout.
`prepare_expert_joint_trajectory`	Prepare a planner result for a destination's fixed command clock.

embodichain.lab.gym.envs.expert_trajectory.EXPERT_TRAJECTORY_SCHEMA_VERSION = 1#

int([x]) -> integer int(x, base=10) -> integer

Convert a number or string to an integer, or return 0 if no arguments are given. If x is a number, return x.__int__(). For floating point numbers, this truncates towards zero.

If x is not a number or if base is given, then x must be a string, bytes, or bytearray instance representing an integer literal in the given base. The literal can be preceded by ‘+’ or ‘-’ and be surrounded by whitespace. The base defaults to 10. Valid bases are 0 and 2-36. Base 0 means to interpret the base from the string as an integer literal. >>> int(‘0b100’, base=0) 4

class embodichain.lab.gym.envs.expert_trajectory.ExpertActionSpec[source]#

Canonical flat action layout used by expert trajectory persistence.

Parameters:

joint_command_mode (Literal['position', 'position_velocity']) – Position-only or position-velocity storage mode.
joint_names (tuple[str, ...]) – Active joint names in stored column order.

Methods:

`__init__`(joint_command_mode, joint_names)
`metadata`(*, step_dt)	Return serializable metadata describing this expert action schema.

Attributes:

`feature_names`	Ordered per-dimension names for dataset feature metadata.
`qpos_slice`	Half-open qpos slice within the stored action vector.
`qvel_slice`	Half-open qvel slice, or `None` for position-only actions.
`width`	Number of scalar values in one stored expert action.

__init__(joint_command_mode, joint_names)#

property feature_names: tuple[str, ...]#: Ordered per-dimension names for dataset feature metadata.

metadata(*, step_dt)[source]#

Return serializable metadata describing this expert action schema.

Parameters:: step_dt (float) – Environment command period associated with stored actions.
Return type:: dict[str, object]
Returns:: JSON-compatible schema metadata.

property qpos_slice: tuple[int, int]#: Half-open qpos slice within the stored action vector.

property qvel_slice: tuple[int, int] | None#: Half-open qvel slice, or None for position-only actions.

property width: int#: Number of scalar values in one stored expert action.

class embodichain.lab.gym.envs.expert_trajectory.ExpertTrajectoryCfg[source]#

Environment-owned expert trajectory execution and recording settings.

Parameters:: joint_command_mode (Literal['position', 'position_velocity']) – Position-only or position-velocity expert targets.

Attributes:

joint_command_mode

Joint targets emitted and stored by expert trajectory generation.

joint_command_mode: Literal['position', 'position_velocity']#: Joint targets emitted and stored by expert trajectory generation.

embodichain.lab.gym.envs.expert_trajectory.JointCommandMode(*args, **kwargs)#: alias of Literal[‘position’, ‘position_velocity’]

embodichain.lab.gym.envs.expert_trajectory.build_expert_action_spec(*, joint_names, joint_command_mode)[source]#

Build the immutable stored-action layout for one environment.

Parameters:

joint_names (Sequence[str]) – Active robot joint names in stored column order.
joint_command_mode (Literal['position', 'position_velocity']) – Position-only or position-velocity storage mode.

Return type:

ExpertActionSpec

Returns:

Validated expert action layout.

Raises:

ValueError – If the mode is unsupported or a joint name is empty.

embodichain.lab.gym.envs.expert_trajectory.encode_expert_action(action, *, spec, active_joint_ids)[source]#

Flatten effective qpos/qvel targets into the expert dataset layout.

Parameters:

action (Tensor | TensorDictBase) – Effective controller action after validation and row masking.
spec (ExpertActionSpec) – Destination expert action layout.
active_joint_ids (Sequence[int]) – Full-robot columns represented by spec.

Return type:

Tensor

Returns:

Active-joint qpos, or concatenated active-joint [qpos, qvel].

Raises:

TypeError – If the action or its joint targets have invalid types.
ValueError – If required targets are missing or cannot be mapped to the configured active joints.

embodichain.lab.gym.envs.expert_trajectory.prepare_expert_joint_trajectory(trajectory, *, control_dt, joint_command_mode)[source]#

Prepare a planner result for a destination’s fixed command clock.

Trajectories require explicit timing and are retimed to control_dt. Velocities are recomputed for position-velocity mode and omitted for position-only mode.

Parameters:

trajectory (PlanResult) – Planner result with positions and explicit timing.
control_dt (float) – Positive destination command period in seconds.
joint_command_mode (Literal['position', 'position_velocity']) – Position-only or position-velocity output mode.

Return type:

PlanResult

Returns:

A detached trajectory aligned with the destination command clock.

Raises:

TypeError – If trajectory has the wrong type.
ValueError – If timing, mode, or required velocity data is invalid.

Controller-ready Actions#

ControllerAction marks controller-ready qpos/qvel/qf commands. The environment keeps the envelope through preprocessing, validates it, and applies it directly without invoking ActionManager. Ordinary flat policy tensors are processed and applied by the manager before physics advances.

class embodichain.lab.gym.envs.types.ControllerAction[source]#

Owned controller-ready action that must still pass through env.step.

The action is already controller-ready. An EmbodiedEnv therefore keeps the envelope through preprocessing, validates the command, and applies it directly without invoking ActionManager.

Parameters:

value (Union[Tensor, TensorDict[str, Tensor]]) – Controller-ready tensor or TensorDict.
metadata (Mapping[str, Any]) – JSON-compatible producer provenance. The environment does not interpret this mapping.

Methods:

`__init__`(value[, metadata])
`snapshot`()	Return an independently owned controller-action envelope.

__init__(value, metadata=<factory>)#

snapshot()[source]#

Return an independently owned controller-action envelope.

Return type:: ControllerAction

Demonstration Episodes#

The segment-aware demonstration API represents a complete task as one episode containing one or more semantic subtasks. Segment action iterables may be lazy, and the common executor records per-environment lengths, terminal status, and segment spans.

Overall language belongs to the task: Task Programs declare a root instruction, and handwritten/direct-control tasks use EmbodiedEnvCfg.task_instruction. An explicit execute_demo_episode instruction overrides these defaults for one episode. The executor freezes the resolved text and its source before planning, so disk recorders and online language batches consume the same semantics even if configuration later changes. Legacy recorder instructions remain a fallback. Segment annotations distinguish explicit labels from overall-task fallback through instruction_source.

class embodichain.lab.gym.envs.demo.DemoExecutionCfg[source]#

Collector-owned settings for demonstration persistence.

segment_fragments persists each eligible program segment as an independent LeRobot episode. It does not resume execution after a failed segment; checkpoint capture and resume are intentionally outside this configuration until an authoritative restore port exists.

Parameters:

mode (Literal['continuous', 'segment_fragments']) – Continuous episode or independent segment-fragment persistence.
save_failed_fragments (bool) – Whether failed segments with recorded frames are retained in fragment mode. Failed fragments remain explicitly annotated and are excluded by successful-segment sampling.

Methods:

`__init__`([mode, save_failed_fragments])
`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.

__init__(mode=<factory>, save_failed_fragments=<factory>)#

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.

embodichain.lab.gym.envs.demo.DemoOutputMode(*args, **kwargs)#

Supported persistence layouts for one demonstration execution.

alias of Literal[‘continuous’, ‘segment_fragments’]

embodichain.lab.gym.envs.demo.DemoSegmentOutcomeKind(*args, **kwargs)#

Stable, first-failure-phase outcome for one program segment row.

alias of Literal[‘succeeded’, ‘runtime_failed’, ‘post_policy_failed’, ‘validation_failed’, ‘cancelled’, ‘truncated’, ‘not_attempted’]

class embodichain.lab.gym.envs.demo.DemoSegment[source]#

One semantic subtask inside a demonstration episode.

The action iterable may be lazy. This lets a task yield one segment, wait for it to execute, inspect the updated scene, and only then plan the next segment.

Parameters:

actions (Iterable[Any]) – Actions for this segment.
name (str) – Stable human-readable segment name.
target_uid (str | None) – Optional scene entity manipulated by this segment.
instruction (str | None) – Optional language instruction specific to this segment.
instruction_source (str | None) – Annotation origin; the executor distinguishes an explicit segment label from a fallback to the overall task.
metadata (Mapping[str, Any]) – Additional JSON-compatible task metadata.
validator (Callable[[], Any] | None) – Optional zero-argument callback that validates this segment after its actions are exhausted. It must return one boolean per parallel environment (or one scalar broadcast to every environment). Gym terminated and truncated remain episode-level signals; use this callback for subtask-level validation.
abort_actions (Callable[..., Iterable[Any]] | None) – Optional callback invoked when the executor stops after retrieving an action but before exhausting the iterable. It receives a reason and last_action_consumed flag, and must return any emergency controller actions that still need ordinary env.step consumption. This is the explicit cancellation handshake for lazy runtimes whose command acknowledgements only mean locally buffered.
failure_policy (Literal['batch_abort', 'row_independent']) – "batch_abort" preserves legacy batch-atomic behavior. "row_independent" permanently freezes only failed environment rows while peers continue through the shared segment and later lazy segments.
progress_total_steps (int | None) – Optional exact action count used by terminal progress wrappers. Leave this as None when the segment can replan, retry, or otherwise emit a data-dependent number of actions.

Methods:

__init__(actions[, name, target_uid, ...])

__init__(actions, name='segment', target_uid=None, instruction=None, metadata=<factory>, validator=None, abort_actions=None, failure_policy='batch_abort', progress_total_steps=None, instruction_source=None)#

class embodichain.lab.gym.envs.demo.DemoSegmentResult[source]#

Execution result and half-open frame range for one segment.

Scalar span and status fields are batch aggregates kept for compatibility. The tuple fields preserve each vector-environment row independently.

Parameters:

segment_id (int) – Zero-based segment index within the episode.
name (str) – Stable segment name supplied by the task.
start_step (int) – Earliest participating row start, inclusive.
end_step (int) – Latest participating row end, exclusive.
success (bool) – Whether every participating row completed the segment.
target_uid (str | None) – Optional manipulated scene entity.
instruction (str | None) – Optional language instruction.
instruction_source (str | None) – Explicit segment label or overall-task fallback.
failure_reason (str | None) – First aggregate failure reason, if any.
metadata (Mapping[str, Any]) – Additional JSON-compatible task metadata.
active (tuple[bool, ...]) – Participation mask captured at segment start.
start_steps (tuple[int, ...]) – Per-environment inclusive starts.
end_steps (tuple[int, ...]) – Per-environment exclusive ends.
successes (tuple[bool, ...]) – Per-environment segment status.
failure_reasons (tuple[str | None, ...]) – Per-environment failure reasons.
attempt_id (int) – Collection attempt that produced this segment.
continuity_id (int) – Causal-continuity region containing this segment.
outcome_kind (Optional[Literal['succeeded', 'runtime_failed', 'post_policy_failed', 'validation_failed', 'cancelled', 'truncated', 'not_attempted']]) – Aggregate first-failure-phase outcome.
outcome_kinds (tuple[Literal['succeeded', 'runtime_failed', 'post_policy_failed', 'validation_failed', 'cancelled', 'truncated', 'not_attempted'], ...]) – Per-environment first-failure-phase outcomes.

Methods:

`__init__`(segment_id, name, start_step, ...)
`to_metadata`([env_id])	Return a JSON-compatible aggregate or per-environment representation.

__init__(segment_id, name, start_step, end_step, success, target_uid=None, instruction=None, failure_reason=None, metadata=<factory>, active=(), start_steps=(), end_steps=(), successes=(), failure_reasons=(), attempt_id=0, continuity_id=0, outcome_kind=None, outcome_kinds=(), instruction_source=None)#

to_metadata(env_id=None)[source]#

Return a JSON-compatible aggregate or per-environment representation.

Parameters:: env_id (int | None) – Optional parallel-environment index. When provided, scalar spans and status are selected from the per-environment fields.
Return type:: dict[str, Any]
Returns:: JSON-compatible segment metadata.

class embodichain.lab.gym.envs.demo.DemoEpisodeResult[source]#

Result of executing all planned segments for one batched episode.

Parameters:

episode_index (int) – Logical episode identifier.
length (int) – Maximum recorded row length.
completed (bool) – Whether every environment completed successfully.
success (tuple[bool, ...]) – Sticky per-environment success flags.
terminated (tuple[bool, ...]) – Sticky per-environment Gym termination flags.
truncated (tuple[bool, ...]) – Sticky per-environment Gym truncation flags.
terminal_reason (str) – Aggregate terminal reason.
segments (tuple[DemoSegmentResult, ...]) – Executed segment results.
lengths (tuple[int, ...]) – Independent per-environment recorded lengths.
completed_by_env (tuple[bool, ...]) – Independent valid-completion flags.
terminal_reasons (tuple[str, ...]) – Independent terminal reasons.
execution_mode (Literal['continuous', 'segment_fragments']) – Persistence layout selected for this execution.
attempt_id (int) – Zero-based collection attempt identifier.
instruction (str | None) – Overall task language frozen before execution.
instruction_source (str | None) – Origin of the overall task language.

Methods:

`__init__`(episode_index, length, completed, ...)
`to_metadata`()	Return a JSON-compatible representation.

Attributes:

`all_success`	Whether every parallel environment completed successfully.
`any_success`	Whether at least one parallel environment completed successfully.
`successful_fragment_count_by_env`	Count accepted, non-empty program segments for each environment.

__init__(episode_index, length, completed, success, terminated, truncated, terminal_reason, segments=(), lengths=(), completed_by_env=(), terminal_reasons=(), execution_mode='continuous', attempt_id=0, instruction=None, instruction_source=None)#

property all_success: bool#: Whether every parallel environment completed successfully.

property any_success: bool#: Whether at least one parallel environment completed successfully.

property successful_fragment_count_by_env: tuple[int, ...]#: Count accepted, non-empty program segments for each environment.

to_metadata()[source]#

Return a JSON-compatible representation.

Return type:: dict[str, Any]

embodichain.lab.gym.envs.demo.execute_demo_episode(env, *, episode_index=0, execution_cfg=None, attempt_id=0, instruction=None, should_stop=None, progress=None, segments=None, **plan_kwargs)[source]#

Plan and execute every segment in one environment episode.

Auto-reset is suspended for the duration of execution. The caller owns the transaction boundary and must explicitly call env.reset() to commit a successful episode or env.reset(options={"save_data": False}) to discard an invalid attempt.

Parameters:

env (Any) – Gym environment or wrapper.
episode_index (int) – Logical episode identifier used in metadata and logs.
execution_cfg (DemoExecutionCfg | None) – Collector-owned output settings. Defaults to continuous episode persistence.
attempt_id (int) – Zero-based identifier for this collection attempt.
instruction (str | None) – Optional overall language override bound to this episode.
should_stop (Callable[[], bool] | None) – Optional callback checked before every action.
progress (Callable[[Iterable[Any], str], Iterable[Any]] | None) – Optional wrapper such as tqdm for action iterables.
segments (Iterable[DemoSegment] | None) – Optional trusted segment iterable. When provided, bypasses task planner selection while retaining ordinary segment validation, instruction fallback, execution, and recording. An empty iterable produces an empty-plan failure. Otherwise planning methods are selected as usual.
**plan_kwargs (Any) – Arguments forwarded to the task’s planning method.

Return type:

DemoEpisodeResult

Returns:

A DemoEpisodeResult describing segment spans and terminal state.

embodichain.lab.gym.envs.demo.resolve_demo_segments(env, **kwargs)[source]#

Resolve a task’s segment plan with legacy single-action-list fallback.

Tasks implementing create_demo_segments own the number, order, and targets of segments. Older tasks that only implement create_demo_action_list are represented as one legacy segment.

Parameters:

env (Any) – Gym environment or wrapper.
**kwargs (Any) – Planning arguments forwarded to the task method.

Return type:

Iterable[DemoSegment]

Returns:

A possibly lazy iterable of DemoSegment objects.

Raises:

AttributeError – If the environment exposes neither planning API.
TypeError – If a segment planner yields a value of the wrong type.

embodichain.lab.gym.envs.demo.resolve_demo_instruction(env, *, instruction=None, task_program=None, legacy_instruction=None)[source]#

Resolve task-owned language independently of recorder construction.

Parameters:

env (Any) – Environment or wrapper exposing task configuration.
instruction (str | None) – Explicit instruction for this episode.
task_program (Any | None) – Actually selected program. Omission uses the static one.
legacy_instruction (Any | None) – Recorder fallback for legacy/external callers.

Returns:

episode, task_program, task, legacy_recorder, or unknown. A selected dynamic program never borrows the static program’s instruction.

Return type:

Overall instruction and its source

Raises:

ValueError – If a provided instruction is invalid or legacy recorders disagree without an authoritative task instruction.

Dynamic Settling#

The shared settling monitor is used by both reset events and Task Program post-policies, so they apply the same row-local stability semantics.

class embodichain.lab.gym.envs.settling.DynamicSettleMonitorCfg[source]#

Threshold and cadence policy for DynamicSettleMonitor.

The monitor never advances an environment. Callers own the stepping path and provide raw velocity samples after the configured minimum/cadence. This lets reset events and demonstration post-policies share exactly the same state transition rules while using different stepping ports.

Methods:

`__init__`([linear_velocity_threshold, ...])
`copy`(**kwargs)	Return a new object replacing specified fields with new values.
`replace`(**kwargs)	Return a new object replacing specified fields with new values.
`snapshot`()	Return an independently owned configuration value.
`to_dict`()	Convert an object into dictionary recursively.
`validate`([prefix])	Check the validity of configclass object.

Attributes:

`angular_velocity_threshold`	Maximum stable angular speed in radians per second.
`check_interval_steps`	Minimum number of steps between independent evidence checks.
`linear_velocity_threshold`	Maximum stable linear speed in metres per second.
`max_steps`	Maximum elapsed environment steps before unresolved rows time out.
`min_steps`	Minimum number of environment steps before the first check.
`required_stable_checks`	Consecutive stable checks required independently for each row.

__init__(linear_velocity_threshold=<factory>, angular_velocity_threshold=<factory>, min_steps=<factory>, max_steps=<factory>, check_interval_steps=<factory>, required_stable_checks=<factory>)#

angular_velocity_threshold: float#: Maximum stable angular speed in radians per second.

check_interval_steps: int#: Minimum number of steps between independent evidence checks.

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.

linear_velocity_threshold: float#: Maximum stable linear speed in metres per second.

max_steps: int#: Maximum elapsed environment steps before unresolved rows time out.

min_steps: int#: Minimum number of environment steps before the first check.

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.

required_stable_checks: int#: Consecutive stable checks required independently for each row.

snapshot()[source]#

Return an independently owned configuration value.

Return type:: DynamicSettleMonitorCfg

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.

class embodichain.lab.gym.envs.settling.DynamicSettleSample[source]#

Raw per-body speed evidence for one registered scene entity.

Parameters:

entity_id (str) – Stable entity identifier used in metadata and diagnostics.
linear_speed (Tensor) – Per-row body speeds with shape (B, N).
angular_speed (Tensor) – Per-row body speeds with shape (B, N).

Methods:

`__init__`(entity_id, linear_speed, angular_speed)
`snapshot`()	Return an independently owned raw evidence sample.

__init__(entity_id, linear_speed, angular_speed)#

snapshot()[source]#

Return an independently owned raw evidence sample.

Return type:: DynamicSettleSample

class embodichain.lab.gym.envs.settling.DynamicSettleState[source]#

Owned state emitted after one monitor observation.

Methods:

`__init__`(env_ids, elapsed_steps, ...)
`to_metadata`()	Return deterministic, JSON-compatible post-policy metadata.

Attributes:

complete

Whether every row has either settled or timed out.

__init__(env_ids, elapsed_steps, observation_count, checked, stable_counts, settled_mask, timeout_mask, max_linear_speed, max_angular_speed)#

property complete: bool#: Whether every row has either settled or timed out.

to_metadata()[source]#

Return deterministic, JSON-compatible post-policy metadata.

Return type:: dict[str, object]

class embodichain.lab.gym.envs.settling.DynamicSettleMonitor[source]#

Track settling independently for stable environment IDs.

Duplicate observations at the same elapsed_steps value are idempotent. Regressing step counters are rejected, and a jump across multiple cadence boundaries counts as one fresh observation rather than replaying one sample.

Methods:

`__init__`(cfg, env_ids)
`observe`(samples, *, elapsed_steps)	Consume one raw speed observation when the configured cadence is due.

Attributes:

env_ids

Return the stable row IDs owned by this monitor.

__init__(cfg, env_ids)[source]#

property env_ids: Tensor#: Return the stable row IDs owned by this monitor.

observe(samples, *, elapsed_steps)[source]#

Consume one raw speed observation when the configured cadence is due.

Parameters:

samples (Sequence[DynamicSettleSample]) – One speed sample per monitored entity.
elapsed_steps (int) – Steps advanced by the caller since post-policy start.

Return type:

DynamicSettleState

Returns:

Per-row stable, settled, timeout, and velocity metadata.

Wrappers#

class embodichain.lab.gym.envs.NoFailWrapper[source]#

Bases: Wrapper

A wrapper that alter the env’s is_task_success method to make sure all the is_task_success determination return True.

Parameters:: env (gym.Env) – the environment to wrap.

Methods:

__init__(env)

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

__init__(env)[source]#

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

Parameters:: env (Env) – The environment to wrap

class embodichain.lab.gym.envs.ReplayWrapper[source]#

Bases: Wrapper

Replay a recorded environment trajectory.

In kinematic mode physics is disabled and every recorded object’s pose/qpos is written directly each step, producing observations only (no reward / success / action). In dynamic mode the recorded robot actions are fed back through env.step() so physics re-simulates the scene; the full obs/reward/terminated/truncated/info tuple is returned. The control mode uses the same kinematic behavior while exposing go_to_step() for interactive scrubbing.

Dynamic replay treats missing meta.action_kind (or "policy") as raw policy actions. Current expert recordings use action_kind="expert"; "expert_controller" remains accepted for stacked-branch compatibility. Expert recordings carry the canonical schema so qpos/qvel targets can be reconstructed before stepping. Schema metadata alone cannot identify controller actions: older position-only raw policy recordings carry it too.

Physical-objective dynamic replay currently supports one environment and full resets only. Stepping past exhaustion raises instead of extending objective dwell time by repeating the final command.

Parameters:

env (Env) – The environment to wrap (constructed without record_trajectory).
trajectory (str | dict) – A .pt path or loaded dict from EmbodiedEnv.save_trajectory().
mode (str) – "kinematic", "dynamic", or "control".

Methods:

`__init__`(env, trajectory[, mode])	Wraps an environment to allow a modular transformation of the `step()` and `reset()` methods.
`close`()	Closes the wrapper and `env`.
`go_to_step`(step)	Scrub to a specific recorded state (kinematic).
`reset`(*[, seed, options])	Uses the `reset()` of the `env` that can be overwritten to change the returned data.
`step`(action)	Uses the `step()` of the `env` that can be overwritten to change the returned data.

Attributes:

control_max_step

Largest state index available to interactive control replay.

__init__(env, trajectory, mode='dynamic')[source]#

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

Parameters:: env (Env) – The environment to wrap

close()[source]#

Closes the wrapper and env.

Return type:: None

property control_max_step: int#: Largest state index available to interactive control replay.

go_to_step(step)[source]#

Scrub to a specific recorded state (kinematic).

State index t is the state immediately before recorded action t.

Parameters:: step (int) – Target step index.
Return type:: TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]
Returns:: The observation at the target step.

reset(*, seed=None, options=None)[source]#

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

Return type:: tuple[TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]], dict]

step(action)[source]#

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

Return type:: tuple[TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]], Tensor, Tensor, Tensor, dict]

`demo`	Segment-aware expert demonstration protocol and executor.
`differentiable_env`	Newton-backed kinematic environment for analytic policy gradient.
`expert_trajectory`	Expert action encoding at the Gym and dataset boundary.
`task_program`	Gym lifecycle bridge for compiled Task Programs.
`managers`	Managers that orchestrate collections of functors (observation, reward, event, action, dataset) running at specific points in the environment step loop.
`types`	Shared action types for the Gym environment boundary.
`wrapper`

embodichain.lab.gym.envs

Contents

embodichain.lab.gym.envs#

Overview#

Environment Classes#

Differentiable Environment#

Expert Trajectories#

Controller-ready Actions#

Demonstration Episodes#

Dynamic Settling#

Wrappers#