embodichain.lab.gym.envs.managers

Contents

embodichain.lab.gym.envs.managers#

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

Each manager owns a typed @configclass config whose attributes are FunctorCfg instances; the config attribute name becomes the functor’s unique identifier.

Overview#

Managers orchestrate collections of functors that run at specific points in the environment step loop. Each manager owns a typed @configclass config whose attributes are FunctorCfg (or subclass) instances; at init the manager resolves every func reference, validates argument signatures against params, resolves SceneEntityCfg targets to scene indices, and groups functors by mode. The config attribute name becomes the functor’s unique identifier within its manager.

The five manager types are ObservationManager (compute(obs)), RewardManager (compute(obs, action, info)), EventManager (apply(mode, env_ids), the home of all randomization functors), ActionManager (process_action(action) then apply_action()), and DatasetManager (step/save for LeRobot recording).

Submodules

randomization

Domain-randomization event functors (physics, visual, spatial, geometry).

Classes

FunctorCfg

Configuration for a functor.

SceneEntityCfg

Configuration for a scene entity that is used by the manager's functor.

EventCfg

Configuration for a event functor.

ObservationCfg

Configuration for an observation functor.

RewardCfg

Configuration for a reward functor.

ActionTermCfg

Configuration for an action term.

DatasetFunctorCfg

Configuration for dataset collection functors.

Functor

Base class for Functor.

ManagerBase

Base class for all managers.

EventManager

Manager for orchestrating operations based on different simulation events.

ObservationManager

Manager for orchestrating operations based on different simulation observations.

RewardManager

Manager for orchestrating reward computation in reinforcement learning tasks.

ActionManager

Split one flat policy action across independently applied terms.

DatasetManager

Manager for orchestrating dataset collection and saving using functors.

ActionTerm

Base class for one independently processed and applied action slice.

ActionDescriptor

Bind one action-term descriptor to a manager-owned flat slice.

ActionTermDescriptor

Describe one term's raw policy-action representation.

ActionTrace

Capture one policy request and its controller-side action stages.

DefaultJointPositionAction

Map normalized actions around a configured default joint pose.

EefPoseAction

Convert an absolute EEF pose into selected-arm joint targets.

JointEffortAction

Apply scaled effort actions to selected joints.

JointPositionAction

Apply absolute position actions to selected joints.

JointPositionToLimitsAction

Map normalized selected-joint actions to position limits.

JointVelocityAction

Apply scaled velocity actions to selected joints.

ParallelGripperAction

Map one scalar policy action to explicit parallel-gripper joints.

RelativeEefPoseAction

Convert normalized arena-frame pose deltas into selected-arm targets.

RelativeJointPositionAction

Add scaled relative actions to current selected positions.

LeRobotRecorder

Functor for recording episodes in LeRobot format.

AsyncLeRobotRecorder

LeRobot recorder that saves episodes on a background thread.

Functions

observations.get_rigid_object_pose(env, obs, ...)

Get the world poses of the rigid objects in the environment.

observations.normalize_robot_joint_data(env, ...)

Normalize selected observation columns using their robot joint limits.

observations.compute_semantic_mask(env, obs, ...)

Compute the semantic mask for the specified scene entity.

observations.compute_exteroception

Compute the exteroception for the observation space.

events.replace_assets_from_group

Replace a directory-sampled mesh or a prepared semantic asset combination.

record.record_camera_data

Record independent environment videos, committed explicitly before reset.

rewards.distance_between_objects(env, obs, ...)

Reward based on distance between two rigid objects.

rewards.success_reward(env, obs, action, info)

Sparse bonus reward when task succeeds.

rewards.distance_to_target(env, obs, action, ...)

Reward based on absolute distance to a virtual target pose.

randomization.visual.randomize_light(env, ...)

Randomize light properties by adding, scaling, or setting random values.

randomization.visual.randomize_camera_intrinsics(...)

Randomize camera intrinsic properties by adding, scaling, or setting random values.

randomization.visual.randomize_visual_material

Randomize the visual material properties of a RigidObject or Articulation.

randomization.spatial.get_random_pose(...[, ...])

Generate a random pose based on the initial position and rotation.

randomization.spatial.randomize_rigid_object_pose(...)

Randomize the pose of a rigid object in the environment.

randomization.spatial.randomize_robot_eef_pose(...)

Randomize the initial end-effector pose of a robot in the environment.

randomization.spatial.randomize_robot_qpos(...)

Randomize the initial joint positions of a robot in the environment.

actions.register_action_contract(...)

Register an implementation for a stable action contract.

actions.resolve_action_contract(contract_id)

Resolve a stable action contract to the current implementation.

register_action_contract(contract_id, ...)

Register an implementation for a stable action contract.

resolve_action_contract(contract_id)

Resolve a stable action contract to the current implementation.

Configuration Classes#

class embodichain.lab.gym.envs.managers.FunctorCfg[source]#

Configuration for a functor.

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:

extra

Extra metadata about the functor.

func

The function or class to be called for the functor.

params

The parameters to be passed to the function as keyword arguments.

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.

extra: dict[str, Any]#

Extra metadata about the functor. Defaults to an empty dict.

This can be used to store additional configuration information such as the output shape of observation functors, which can be used for pre-allocating buffers.

For observation functors, common keys include:
  • shape: A tuple defining the output shape of the functor (excluding num_envs dimension).

func: Callable | Functor#

The function or class to be called for the functor.

The function must take the environment object as the first argument. The remaining arguments are specified in the params attribute.

It also supports callable classes, i.e. classes that implement the __call__() method. In this case, the class should inherit from the Functor class and implement the required methods.

params: dict[str, Any | SceneEntityCfg]#

The parameters to be passed to the function as keyword arguments. Defaults to an empty dict.

Note

If the value is a SceneEntityCfg object, the manager will query the scene entity from the SimulationManager and process the entity’s joints and bodies as specified in the SceneEntityCfg 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.

class embodichain.lab.gym.envs.managers.SceneEntityCfg[source]#

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

This class is used to specify the name of the scene entity that is queried from the SimulationManager and passed to the manager’s functor.

Attributes:

body_ids

The indices of the bodies from the asset required by the functor.

body_names

The names of the bodies from the asset required by the functor.

control_parts

The names of the control parts from the asset(only support for robot) required by the functor.

joint_ids

The indices of the joints from the asset required by the functor.

joint_names

The names of the joints from the scene entity.

link_names

The names of the links from the asset required by the functor.

preserve_order

Whether to preserve indices ordering to match with that in the specified joint, body, or object collection names.

uid

The name of the scene entity.

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.

resolve(scene)

Resolves the scene entity and converts the joint and body names to indices.

to_dict()

Convert an object into dictionary recursively.

validate([prefix])

Check the validity of configclass object.

body_ids: list[int] | slice#

The indices of the bodies from the asset required by the functor. Defaults to slice(None), which means all the bodies in the asset.

If body_names is specified, this is filled in automatically on initialization of the manager.

body_names: str | list[str] | None#

The names of the bodies from the asset required by the functor. Defaults to None.

The names can be either body names or a regular expression matching the body names.

These are converted to body indices on initialization of the manager and passed to the functor function as a list of body indices under body_ids.

control_parts: str | list[str] | None#

The names of the control parts from the asset(only support for robot) required by the functor. Defaults to None.

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.

joint_ids: list[int] | slice#

The indices of the joints from the asset required by the functor. Defaults to slice(None), which means all the joints in the asset (if present).

If joint_names is specified, this is filled in automatically on initialization of the manager.

joint_names: str | list[str] | None#

The names of the joints from the scene entity. Defaults to None.

The names can be either joint names or a regular expression matching the joint names.

These are converted to joint indices on initialization of the manager and passed to the functor as a list of joint indices under joint_ids.

The names of the links from the asset required by the functor. Defaults to None.

The names can be either link names or a regular expression matching the link names.

preserve_order: bool#

Whether to preserve indices ordering to match with that in the specified joint, body, or object collection names. Defaults to False.

If False, the ordering of the indices are sorted in ascending order (i.e. the ordering in the entity’s joints, bodies, or object in the object collection). Otherwise, the indices are preserved in the order of the specified joint, body, or object collection names.

For more details, see the isaaclab.utils.string.resolve_matching_names() function.

Note

This attribute is only used when joint_names, body_names are specified.

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.

resolve(scene)[source]#

Resolves the scene entity and converts the joint and body names to indices.

This function examines the scene entity from the SimulationManager and resolves the indices and names of the joints and bodies. It is an expensive operation as it resolves regular expressions and should be called only once.

Parameters:

scene (SimulationManager) – The interactive scene instance.

Raises:
  • ValueError – If the scene entity is not found.

  • ValueError – If both joint_names and joint_ids are specified and are not consistent.

  • ValueError – If both body_names and body_ids are specified and are not consistent.

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.

uid: str#

The name of the scene entity.

This is the name defined in the scene configuration file. See the SimulationManagerCfg class for more details.

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

Configuration for a event functor.

The event functor is used to trigger events in the environment at specific times or under specific conditions. The mode attribute determines when the functor is applied. - startup: The functor is applied when the environment is started. - interval: The functor is applied at each env step. - reset: The functor is applied when the environment is reset.

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:

extra

Extra metadata about the functor.

func

The function or class to be called for the functor.

interval_step

The number of environment step after which the functor is applied.

is_global

Whether the event should be tracked on a per-environment basis.

mode

The mode in which the event functor is applied.

params

The parameters to be passed to the function as keyword arguments.

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.

extra: dict[str, Any]#

Extra metadata about the functor. Defaults to an empty dict.

This can be used to store additional configuration information such as the output shape of observation functors, which can be used for pre-allocating buffers.

For observation functors, common keys include:
  • shape: A tuple defining the output shape of the functor (excluding num_envs dimension).

func: Callable | Functor#

The function or class to be called for the functor.

The function must take the environment object as the first argument. The remaining arguments are specified in the params attribute.

It also supports callable classes, i.e. classes that implement the __call__() method. In this case, the class should inherit from the Functor class and implement the required methods.

interval_step: int#

The number of environment step after which the functor is applied. Defaults to 4.

is_global: bool#

Whether the event should be tracked on a per-environment basis. Defaults to False.

If True, the same interval step is used for all the environment instances. If False, the interval step is sampled independently for each environment instance and the functor is applied when the current step hits the interval step for that instance.

Note

This is only used if the mode is "interval".

mode: Literal['startup', 'interval', 'reset']#

The mode in which the event functor is applied.

Note

The mode name "interval" is a special mode that is handled by the manager Hence, its name is reserved and cannot be used for other modes.

params: dict[str, Any | SceneEntityCfg]#

The parameters to be passed to the function as keyword arguments. Defaults to an empty dict.

Note

If the value is a SceneEntityCfg object, the manager will query the scene entity from the SimulationManager and process the entity’s joints and bodies as specified in the SceneEntityCfg 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.

class embodichain.lab.gym.envs.managers.ObservationCfg[source]#

Configuration for an observation functor.

The observation functor is used to compute observations for the environment. The mode attribute determines whether the observation is already present in the observation space or not.

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:

extra

Extra metadata about the functor.

func

The function or class to be called for the functor.

mode

The mode for the observation computation.

name

The name of the observation.

params

The parameters to be passed to the function as keyword arguments.

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.

extra: dict[str, Any]#

Extra metadata about the functor. Defaults to an empty dict.

This can be used to store additional configuration information such as the output shape of observation functors, which can be used for pre-allocating buffers.

For observation functors, common keys include:
  • shape: A tuple defining the output shape of the functor (excluding num_envs dimension).

func: Callable | Functor#

The function or class to be called for the functor.

The function must take the environment object as the first argument. The remaining arguments are specified in the params attribute.

It also supports callable classes, i.e. classes that implement the __call__() method. In this case, the class should inherit from the Functor class and implement the required methods.

mode: Literal['modify', 'add']#

The mode for the observation computation.

  • modify: The observation is already present in the observation space, updated the value in-place.

  • add: The observation is not present in the observation space, add a new entry to the observation space.

name: str#

The name of the observation.

The name can be a new key to observation space, eg:
  • object_position: shape of (num_envs, 3)

  • robot/eef_pose: shape of (num_envs, 7) or (num_envs, 4, 4)

  • sensor/cam_high/mask: shape of (num_envs, H, W)

or a existing key to modify, eg:
  • robot/qpos: shape of (num_envs, num_dofs)

/ is used to separate different levels of hierarchy in the observation dictionary.

params: dict[str, Any | SceneEntityCfg]#

The parameters to be passed to the function as keyword arguments. Defaults to an empty dict.

Note

If the value is a SceneEntityCfg object, the manager will query the scene entity from the SimulationManager and process the entity’s joints and bodies as specified in the SceneEntityCfg 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.

class embodichain.lab.gym.envs.managers.RewardCfg[source]#

Configuration for a reward functor.

The reward functor is used to compute rewards for the environment. The mode attribute determines how the reward is combined with existing rewards.

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:

extra

Extra metadata about the functor.

func

The function or class to be called for the functor.

mode

The mode for the reward computation.

params

The parameters to be passed to the function as keyword arguments.

weight

The weight multiplier for this reward term.

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.

extra: dict[str, Any]#

Extra metadata about the functor. Defaults to an empty dict.

This can be used to store additional configuration information such as the output shape of observation functors, which can be used for pre-allocating buffers.

For observation functors, common keys include:
  • shape: A tuple defining the output shape of the functor (excluding num_envs dimension).

func: Callable | Functor#

The function or class to be called for the functor.

The function must take the environment object as the first argument. The remaining arguments are specified in the params attribute.

It also supports callable classes, i.e. classes that implement the __call__() method. In this case, the class should inherit from the Functor class and implement the required methods.

mode: Literal['add', 'replace']#

The mode for the reward computation.

  • add: The reward is added to the existing total reward.

  • replace: The reward replaces the total reward (useful for single reward functions).

params: dict[str, Any | SceneEntityCfg]#

The parameters to be passed to the function as keyword arguments. Defaults to an empty dict.

Note

If the value is a SceneEntityCfg object, the manager will query the scene entity from the SimulationManager and process the entity’s joints and bodies as specified in the SceneEntityCfg 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.

weight: float#

The weight multiplier for this reward term.

This value is used to scale the reward before adding it to the total reward. Default is 1.0 (no scaling).

class embodichain.lab.gym.envs.managers.ActionTermCfg[source]#

Configuration for an action term.

The action term owns one ordered slice of the flat policy action. New persisted configurations should use contract so the action semantics remain stable when the implementation class is renamed. A configured contract must match the selected action implementation’s contract_id.

Attributes:

contract

Stable semantic contract resolved to a current action implementation.

extra

Extra metadata about the functor.

func

The function or class to be called for the functor.

params

The parameters to be passed to the function as keyword arguments.

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.

contract: str | None#

Stable semantic contract resolved to a current action implementation.

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.

extra: dict[str, Any]#

Extra metadata about the functor. Defaults to an empty dict.

This can be used to store additional configuration information such as the output shape of observation functors, which can be used for pre-allocating buffers.

For observation functors, common keys include:
  • shape: A tuple defining the output shape of the functor (excluding num_envs dimension).

func: Callable | Functor#

The function or class to be called for the functor.

The function must take the environment object as the first argument. The remaining arguments are specified in the params attribute.

It also supports callable classes, i.e. classes that implement the __call__() method. In this case, the class should inherit from the Functor class and implement the required methods.

params: dict[str, Any | SceneEntityCfg]#

The parameters to be passed to the function as keyword arguments. Defaults to an empty dict.

Note

If the value is a SceneEntityCfg object, the manager will query the scene entity from the SimulationManager and process the entity’s joints and bodies as specified in the SceneEntityCfg 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.

class embodichain.lab.gym.envs.managers.DatasetFunctorCfg[source]#

Configuration for dataset collection functors.

Dataset functors are called with mode=”save” which handles both: - Recording observation-action pairs on every step - Auto-saving episodes when dones=True

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:

extra

Extra metadata about the functor.

func

The function or class to be called for the functor.

params

The parameters to be passed to the function as keyword arguments.

save_failed_episodes

Whether to save failed episodes.

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.

extra: dict[str, Any]#

Extra metadata about the functor. Defaults to an empty dict.

This can be used to store additional configuration information such as the output shape of observation functors, which can be used for pre-allocating buffers.

For observation functors, common keys include:
  • shape: A tuple defining the output shape of the functor (excluding num_envs dimension).

func: Callable | Functor#

The function or class to be called for the functor.

The function must take the environment object as the first argument. The remaining arguments are specified in the params attribute.

It also supports callable classes, i.e. classes that implement the __call__() method. In this case, the class should inherit from the Functor class and implement the required methods.

params: dict[str, Any | SceneEntityCfg]#

The parameters to be passed to the function as keyword arguments. Defaults to an empty dict.

Note

If the value is a SceneEntityCfg object, the manager will query the scene entity from the SimulationManager and process the entity’s joints and bodies as specified in the SceneEntityCfg 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.

save_failed_episodes: bool#

Whether to save failed episodes.

If enabled for any dataset functor, all save-mode dataset functors receive both successful and failed episodes. During run-task expert generation, a non-empty failed or truncated result is committed and counts toward max_episodes instead of being discarded and retried. Empty plans and execution exceptions remain uncommitted.

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.

Base Classes#

class embodichain.lab.gym.envs.managers.Functor[source]#

Bases: ABC

Base class for Functor.

Functor implementations can be functions or classes. If the functor is a class, it should inherit from this base class and implement the required methods.

Each manager is implemented as a class that inherits from the ManagerBase class. Each manager class should also have a corresponding configuration class that defines the configuration functors for the manager. Each functor should the FunctorCfg class or its subclass.

Example pseudo-code for creating a manager:

from embodichain.utils import configclass
from embodichain.lab.gym.managers import ManagerBase
from embodichain.lab.gym.managers FunctorCfg

@configclass
class MyManagerCfg:

    functor1: FunctorCfg = FunctorCfg(...)
    functor2: FunctorCfg = FunctorCfg(...)
    functor3: FunctorCfg = FunctorCfg(...)

# define manager instance
my_manager = ManagerBase(cfg=ManagerCfg(), env=env)

Methods:

__init__(cfg, env)

Initialize the functor.

reset([env_ids])

Resets the functor.

serialize()

General serialization call.

Attributes:

device

Device on which to perform computations.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

__init__(cfg, env)[source]#

Initialize the functor.

Parameters:
property device: str#

Device on which to perform computations.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

reset(env_ids=None)[source]#

Resets the functor.

Parameters:

env_ids (Optional[Sequence[int]]) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

serialize()[source]#

General serialization call. Includes the configuration dict.

Return type:

dict

class embodichain.lab.gym.envs.managers.ManagerBase[source]#

Bases: ABC

Base class for all managers.

Methods:

__init__(cfg, env)

Initialize the manager.

find_functors(name_keys)

Find functors in the manager based on the names.

get_active_iterable_functors(env_idx)

Returns the active functors as iterable sequence of tuples.

reset([env_ids])

Resets the manager and returns logging information for the current time-step.

Attributes:

active_functors

Name of active functors.

device

Device on which to perform computations.

num_envs

Number of environments.

__init__(cfg, env)[source]#

Initialize the manager.

This function is responsible for parsing the configuration object and creating the functors.

If the simulation is not playing, the scene entities are not resolved immediately. Instead, the resolution is deferred until the simulation starts. This is done to ensure that the scene entities are resolved even if the manager is created after the simulation has already started.

Parameters:
  • cfg (object) – The configuration object. If None, the manager is initialized without any functors.

  • env (EmbodiedEnv) – The environment instance.

abstract property active_functors: list[str] | dict[str, list[str]]#

Name of active functors.

property device: str#

Device on which to perform computations.

find_functors(name_keys)[source]#

Find functors in the manager based on the names.

This function searches the manager for functors based on the names. The names can be specified as regular expressions or a list of regular expressions. The search is performed on the active functors in the manager.

Please check the resolve_matching_names() function for more information on the name matching.

Parameters:

name_keys (Union[str, Sequence[str]]) – A regular expression or a list of regular expressions to match the functor names.

Return type:

list[str]

Returns:

A list of functor names that match the input keys.

get_active_iterable_functors(env_idx)[source]#

Returns the active functors as iterable sequence of tuples.

The first element of the tuple is the name of the functor and the second element is the raw value(s) of the functor.

Return type:

Sequence[tuple[str, Sequence[float]]]

Returns:

The active functors.

property num_envs: int#

Number of environments.

reset(env_ids=None)[source]#

Resets the manager and returns logging information for the current time-step.

Parameters:

env_ids (Optional[Sequence[int]]) – The environment ids for which to log data. Defaults None, which logs data for all environments.

Return type:

dict[str, float]

Returns:

Dictionary containing the logging information.

Managers#

class embodichain.lab.gym.envs.managers.EventManager[source]#

Bases: ManagerBase

Manager for orchestrating operations based on different simulation events.

The event manager applies operations to the environment based on different simulation events. For example, changing the masses of objects or their friction coefficients during initialization/ reset, or applying random pushes to the robot at a fixed interval of steps. The user can specify several modes of events to fine-tune the behavior based on when to apply the event.

The event functors are parsed from a config class containing the manager’s settings and each functor’s parameters. Each event functor should instantiate the EventCfg class.

Event functors can be grouped by their mode. The mode is a user-defined string that specifies when the event functor should be applied. This provides the user complete control over when event functors should be applied.

For a typical training process, you may want to apply events in the following modes:

  • “prestartup”: Event is applied once at the beginning of the training before the simulation starts. This is used to randomize USD-level properties of the simulation stage.

  • “startup”: Event is applied once at the beginning of the training once simulation is started.

  • “reset”: Event is applied at every reset.

  • “interval”: Event is applied at pre-specified intervals of time.

However, you can also define your own modes and use them in the training process as you see fit. For this you will need to add the triggering of that mode in the environment implementation as well.

Note

The triggering of operations corresponding to the mode "interval" are the only mode that are directly handled by the manager itself. The other modes are handled by the environment implementation.

Methods:

__init__(cfg, env)

Initialize the event manager.

apply(mode[, env_ids])

Calls each event functor in the specified mode.

find_functors(name_keys)

Find functors in the manager based on the names.

get_active_iterable_functors(env_idx)

Returns the active functors as iterable sequence of tuples.

get_functor(functor_name)

Retrieve a functor from the configuration by its name.

get_functor_cfg(functor_name)

Gets the configuration for the specified functor.

reset([env_ids])

Resets the manager and returns logging information for the current time-step.

set_functor_cfg(functor_name, cfg)

Sets the configuration of the specified functor into the manager.

set_seed(seed)

Set the event seed and rewind all deterministic event streams.

Attributes:

active_functors

Name of active event functors.

available_modes

Modes of events.

device

Device on which to perform computations.

num_envs

Number of environments.

seed

Return the base seed used for event-functor randomization.

seed_revision

Return the explicit reseed revision, including repeated seed values.

__init__(cfg, env)[source]#

Initialize the event manager.

Parameters:
  • cfg (object) – A configuration object or dictionary (dict[str, EventCfg]).

  • env (EmbodiedEnv) – An environment object.

property active_functors: dict[str, list[str]]#

Name of active event functors.

The keys are the modes of event and the values are the names of the event functors.

apply(mode, env_ids=None)[source]#

Calls each event functor in the specified mode.

This function iterates over all the event functors in the specified mode and calls the function corresponding to the functor. The function is called with the environment instance and the environment indices to apply the event to.

For the “interval” mode, the function is called when the time interval has passed. This requires specifying the time step of the environment.

For the “reset” mode, the function is called when the mode is “reset” and the total number of environment steps that have happened since the last trigger of the function is equal to its configured parameter for the number of environment steps between resets.

Parameters:
  • mode (str) – The mode of event.

  • env_ids (Optional[Sequence[int]]) – The indices of the environments to apply the event to. Defaults to None, in which case the event is applied to all environments when applicable.

Raises:
  • ValueError – If the mode is "interval" and the environment indices are provided. This is an undefined behavior as the environment indices are computed based on the time left for each environment.

  • ValueError – If the mode is "reset" and the total number of environment steps that have happened is not provided.

property available_modes: list[str]#

Modes of events.

property device: str#

Device on which to perform computations.

find_functors(name_keys)#

Find functors in the manager based on the names.

This function searches the manager for functors based on the names. The names can be specified as regular expressions or a list of regular expressions. The search is performed on the active functors in the manager.

Please check the resolve_matching_names() function for more information on the name matching.

Parameters:

name_keys (Union[str, Sequence[str]]) – A regular expression or a list of regular expressions to match the functor names.

Return type:

list[str]

Returns:

A list of functor names that match the input keys.

get_active_iterable_functors(env_idx)#

Returns the active functors as iterable sequence of tuples.

The first element of the tuple is the name of the functor and the second element is the raw value(s) of the functor.

Return type:

Sequence[tuple[str, Sequence[float]]]

Returns:

The active functors.

get_functor(functor_name)[source]#

Retrieve a functor from the configuration by its name.

Parameters:

functor_name (str) – The name of the functor to retrieve.

Returns:

The functor if it exists in the configuration, otherwise None.

get_functor_cfg(functor_name)[source]#

Gets the configuration for the specified functor.

The method finds the functor by name by searching through all the modes. It then returns the configuration of the functor with the first matching name.

Parameters:

functor_name (str) – The name of the event functor.

Return type:

EventCfg

Returns:

The configuration of the event functor.

Raises:

ValueError – If the functor name is not found.

property num_envs: int#

Number of environments.

reset(env_ids=None)[source]#

Resets the manager and returns logging information for the current time-step.

Parameters:

env_ids (Optional[Sequence[int]]) – The environment ids for which to log data. Defaults None, which logs data for all environments.

Return type:

dict[str, float]

Returns:

Dictionary containing the logging information.

property seed: int | None#

Return the base seed used for event-functor randomization.

property seed_revision: int#

Return the explicit reseed revision, including repeated seed values.

set_functor_cfg(functor_name, cfg)[source]#

Sets the configuration of the specified functor into the manager.

The method finds the functor by name by searching through all the modes. It then updates the configuration of the functor with the first matching name.

Parameters:
  • functor_name (str) – The name of the event functor.

  • cfg (EventCfg) – The configuration for the event functor.

Raises:

ValueError – If the functor name is not found.

set_seed(seed)[source]#

Set the event seed and rewind all deterministic event streams.

Parameters:

seed (int | None) – Base task-environment seed. None disables scoped event randomization and preserves the process-global RNG behavior.

Raises:

TypeError – If seed is not an integer or None.

Return type:

None

class embodichain.lab.gym.envs.managers.ObservationManager[source]#

Bases: ManagerBase

Manager for orchestrating operations based on different simulation observations.

The default observation space will contain two observation groups:
  • robot: Contains the default observations related to the robot.
    • qpos: The joint positions of the robot.

    • qvel: The joint velocities of the robot.

    • qf: The joint forces of the robot.

  • sensor: Contains the observations related to the sensors which are enabled in the environment.

The observation manager offers two modes of operation:
  • modify: This mode perform data fetching and modification on existing observation data.

  • add: This mode perform new observation computation and add new observation data to the observation space.

Methods:

__init__(cfg, env)

Initialize the observation manager.

compute(obs)

Calls each observation functor in the specified mode.

find_functors(name_keys)

Find functors in the manager based on the names.

get_active_iterable_functors(env_idx)

Returns the active functors as iterable sequence of tuples.

get_functor_cfg(functor_name)

Gets the configuration for the specified functor.

reset([env_ids])

Resets the manager and returns logging information for the current time-step.

Attributes:

active_functors

Name of active observation functors.

device

Device on which to perform computations.

num_envs

Number of environments.

__init__(cfg, env)[source]#

Initialize the observation manager.

Parameters:
  • cfg (object) – A configuration object or dictionary (dict[str, ObservationCfg]).

  • env (EmbodiedEnv) – An environment object.

property active_functors: dict[str, list[str]]#

Name of active observation functors.

The keys are the modes of observation and the values are the names of the observation functors.

compute(obs)[source]#

Calls each observation functor in the specified mode.

This function iterates over all the observation functors in the specified mode and calls the function corresponding to the functor. The function is called with the environment instance and the environment indices to apply the observation to.

Parameters:

obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation data to apply the observation to.

Return type:

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

Returns:

The modified observation data.

Raises:

ValueError – If the mode is not supported.

property device: str#

Device on which to perform computations.

find_functors(name_keys)#

Find functors in the manager based on the names.

This function searches the manager for functors based on the names. The names can be specified as regular expressions or a list of regular expressions. The search is performed on the active functors in the manager.

Please check the resolve_matching_names() function for more information on the name matching.

Parameters:

name_keys (Union[str, Sequence[str]]) – A regular expression or a list of regular expressions to match the functor names.

Return type:

list[str]

Returns:

A list of functor names that match the input keys.

get_active_iterable_functors(env_idx)#

Returns the active functors as iterable sequence of tuples.

The first element of the tuple is the name of the functor and the second element is the raw value(s) of the functor.

Return type:

Sequence[tuple[str, Sequence[float]]]

Returns:

The active functors.

get_functor_cfg(functor_name)[source]#

Gets the configuration for the specified functor.

The method finds the functor by name by searching through all the modes. It then returns the configuration of the functor with the first matching name.

Parameters:

functor_name (str) – The name of the observation functor.

Return type:

ObservationCfg

Returns:

The configuration of the observation functor.

Raises:

ValueError – If the functor name is not found.

property num_envs: int#

Number of environments.

reset(env_ids=None)[source]#

Resets the manager and returns logging information for the current time-step.

Parameters:

env_ids (Optional[Sequence[int]]) – The environment ids for which to log data. Defaults None, which logs data for all environments.

Return type:

dict[str, float]

Returns:

Dictionary containing the logging information.

class embodichain.lab.gym.envs.managers.RewardManager[source]#

Bases: ManagerBase

Manager for orchestrating reward computation in reinforcement learning tasks.

The reward manager computes rewards based on the current state of the environment and actions. It supports multiple reward terms that can be combined through weighted summation.

The reward manager offers two modes of operation:
  • add: This mode computes a reward term and adds it to the total reward (weighted by the term’s weight).

  • replace: This mode replaces the total reward with the computed value (useful for single reward functions).

Note: The config key is used as the unique identifier and display name for each reward functor.

Methods:

__init__(cfg, env)

Initialize the reward manager.

compute(obs, action, info)

Compute the total reward by calling each reward functor.

find_functors(name_keys)

Find functors in the manager based on the names.

get_active_iterable_functors(env_idx)

Returns the active functors as iterable sequence of tuples.

get_functor_cfg(functor_name)

Gets the configuration for the specified functor.

reset([env_ids])

Reset reward terms that are stateful (implemented as classes).

Attributes:

active_functors

Name of active reward functors.

device

Device on which to perform computations.

num_envs

Number of environments.

__init__(cfg, env)[source]#

Initialize the reward manager.

Parameters:
  • cfg (object) – A configuration object or dictionary (dict[str, RewardCfg]).

  • env (EmbodiedEnv) – An environment object.

property active_functors: dict[str, list[str]]#

Name of active reward functors.

The keys are the modes of reward computation and the values are the names of the reward functors.

compute(obs, action, info)[source]#

Compute the total reward by calling each reward functor.

This function iterates over all the reward functors and calls them to compute individual reward terms. The terms are then combined according to their mode and weight.

Parameters:
  • obs (EnvObs) – The observation from the environment.

  • action (EnvAction) – The action applied to the robot.

  • info (dict) – Additional information dictionary.

Returns:

  • total_reward: The total reward for each environment (shape: [num_envs]).

  • reward_info: A dictionary mapping reward term names to their values for logging.

Return type:

A tuple containing

Raises:

ValueError – If the mode is not supported.

property device: str#

Device on which to perform computations.

find_functors(name_keys)#

Find functors in the manager based on the names.

This function searches the manager for functors based on the names. The names can be specified as regular expressions or a list of regular expressions. The search is performed on the active functors in the manager.

Please check the resolve_matching_names() function for more information on the name matching.

Parameters:

name_keys (Union[str, Sequence[str]]) – A regular expression or a list of regular expressions to match the functor names.

Return type:

list[str]

Returns:

A list of functor names that match the input keys.

get_active_iterable_functors(env_idx)#

Returns the active functors as iterable sequence of tuples.

The first element of the tuple is the name of the functor and the second element is the raw value(s) of the functor.

Return type:

Sequence[tuple[str, Sequence[float]]]

Returns:

The active functors.

get_functor_cfg(functor_name)[source]#

Gets the configuration for the specified functor.

The method finds the functor by name by searching through all the modes. It then returns the configuration of the functor with the first matching name.

Parameters:

functor_name (str) – The name of the reward functor.

Return type:

RewardCfg

Returns:

The configuration of the reward functor.

Raises:

ValueError – If the functor name is not found.

property num_envs: int#

Number of environments.

reset(env_ids=None)[source]#

Reset reward terms that are stateful (implemented as classes).

Parameters:

env_ids (Optional[Sequence[int]]) – The environment indices to reset. If None, all environments are reset.

Return type:

dict[str, float]

Returns:

An empty dictionary (no logging needed for reset).

class embodichain.lab.gym.envs.managers.ActionManager[source]#

Bases: ManagerBase

Split one flat policy action across independently applied terms.

Methods:

__init__(cfg, env)

Initialize action terms, their slices, and manager-owned history.

action_trace([executed])

Return a snapshot of requested, processed, and executed commands.

apply_action()

Apply every processed term command in configuration order.

find_functors(name_keys)

Find functors in the manager based on the names.

get_active_iterable_functors(env_idx)

Returns the active functors as iterable sequence of tuples.

get_term(name)

Return one configured term by name.

get_term_by_contract(contract_id)

Return the unique action term implementing a semantic contract.

mask_inactive(active_mask)

Replace inactive vector rows with resource-safe commands.

process_action(action)

Validate, store, and distribute one complete flat policy action.

reset([env_ids])

Reset manager and term state for selected environments.

Attributes:

action

Return current complete raw policy actions.

active_functors

Return active term names in flat policy order.

descriptors

Return ordered action-term descriptors bound to flat slices.

device

Device on which to perform computations.

num_envs

Number of environments.

previous_action

Return previous complete raw policy actions.

single_action_space

Return concatenated raw action bounds in term order.

total_action_dim

Return the width of the complete flat policy action.

__init__(cfg, env)[source]#

Initialize action terms, their slices, and manager-owned history.

Parameters:
  • cfg (object) – Mapping or component object containing ActionTermCfg values.

  • env (EmbodiedEnv) – Owning embodied environment.

property action: Tensor#

Return current complete raw policy actions.

action_trace(executed=None)[source]#

Return a snapshot of requested, processed, and executed commands.

Parameters:

executed (dict[str, Tensor] | None) – Optional term-keyed command snapshot captured at the controller boundary. When omitted, the processed term values are used as the best available command representation.

Return type:

ActionTrace

Returns:

An independently owned ActionTrace.

property active_functors: list[str]#

Return active term names in flat policy order.

apply_action()[source]#

Apply every processed term command in configuration order.

Return type:

None

property descriptors: tuple[ActionDescriptor, ...]#

Return ordered action-term descriptors bound to flat slices.

property device: str#

Device on which to perform computations.

find_functors(name_keys)#

Find functors in the manager based on the names.

This function searches the manager for functors based on the names. The names can be specified as regular expressions or a list of regular expressions. The search is performed on the active functors in the manager.

Please check the resolve_matching_names() function for more information on the name matching.

Parameters:

name_keys (Union[str, Sequence[str]]) – A regular expression or a list of regular expressions to match the functor names.

Return type:

list[str]

Returns:

A list of functor names that match the input keys.

get_active_iterable_functors(env_idx)#

Returns the active functors as iterable sequence of tuples.

The first element of the tuple is the name of the functor and the second element is the raw value(s) of the functor.

Return type:

Sequence[tuple[str, Sequence[float]]]

Returns:

The active functors.

get_term(name)[source]#

Return one configured term by name.

Return type:

ActionTerm

get_term_by_contract(contract_id)[source]#

Return the unique action term implementing a semantic contract.

Parameters:

contract_id (str) – Versioned implementation-independent contract ID.

Raises:
  • KeyError – If no term implements contract_id.

  • ValueError – If multiple terms implement the same contract.

Return type:

ActionTerm

mask_inactive(active_mask)[source]#

Replace inactive vector rows with resource-safe commands.

Position terms hold measured selected-joint positions. Velocity and effort terms use zero commands. This is used by sticky vectorized demo execution after term processing and before application.

Parameters:

active_mask (Tensor) – Boolean tensor of shape (num_envs,) on the environment device. True rows retain processed commands.

Raises:
  • TypeError – If the mask is not boolean.

  • ValueError – If mask or processed-command layout is invalid.

  • RuntimeError – If a resource type has no defined safe inactive value.

Return type:

None

property num_envs: int#

Number of environments.

property previous_action: Tensor#

Return previous complete raw policy actions.

process_action(action)[source]#

Validate, store, and distribute one complete flat policy action.

Parameters:

action (Tensor) – Floating tensor with shape (num_envs, total_action_dim).

Raises:
  • TypeError – If the action is not a floating tensor.

  • ValueError – If shape or device does not match the environment.

Return type:

None

reset(env_ids=None)[source]#

Reset manager and term state for selected environments.

Parameters:

env_ids (list[int] | Tensor | None) – Selected rows. None resets every row.

Return type:

dict[str, float]

Returns:

Empty diagnostic mapping.

property single_action_space: Box#

Return concatenated raw action bounds in term order.

property total_action_dim: int#

Return the width of the complete flat policy action.

class embodichain.lab.gym.envs.managers.DatasetManager[source]#

Bases: ManagerBase

Manager for orchestrating dataset collection and saving using functors.

The dataset manager supports multiple dataset formats through a functor system: - LeRobot format (via LeRobotRecorder) - HDF5 format (via HDF5Recorder) - Zarr format (via ZarrRecorder) - Custom formats (via user-defined functors)

Each functor’s step() method is called once per environment step and handles: - Recording observation-action pairs - Detecting episode completion (dones=True) - Auto-saving completed episodes

Example configuration:
>>> from embodichain.lab.gym.envs.managers.cfg import DatasetFunctorCfg
>>> from embodichain.lab.gym.envs.managers.datasets import LeRobotRecorder
>>>
>>> @configclass
>>> class MyEnvCfg:
>>>     dataset: dict = {
>>>         "lerobot": DatasetFunctorCfg(
>>>             func=LeRobotRecorder,
>>>             save_failed_episodes=True,
>>>             params={
>>>                 "robot_meta": {...},
>>>                 "instruction": {"lang": "pick and place"},
>>>                 "extra": {"scene_type": "kitchen"},
>>>                 "save_path": "/data/datasets"
>>>             }
>>>         )
>>>     }

Methods:

__init__(cfg, env)

Initialize the dataset manager.

apply(mode[, env_ids])

Apply dataset functors for the specified mode.

clear_cache()

Clear cached data from all dataset functors (for online training).

close()

Finalize all dataset functors; repeated calls are safe.

finalize()

Finalize every dataset functor exactly once.

find_functors(name_keys)

Find functors in the manager based on the names.

get_active_iterable_functors(env_idx)

Returns the active functors as iterable sequence of tuples.

get_cached_data()

Get cached data from all dataset functors (for online training).

get_functor_cfg(functor_name)

Gets the configuration for the specified functor.

reset([env_ids])

Reset all dataset functors.

validate_policy_action(action)

Validate a live policy action for every descriptor-driven recorder.

Attributes:

active_functors

Name of active dataset functors by mode.

available_modes

List of available modes for the dataset manager.

device

Device on which to perform computations.

num_envs

Number of environments.

requires_raw_actions

Whether a recorder needs pre-controller policy actions.

save_failed_episodes

Whether any configured dataset recorder should keep failed episodes.

__init__(cfg, env)[source]#

Initialize the dataset manager.

Parameters:
  • cfg (object) – Configuration object containing dataset functor configurations.

  • env (EmbodiedEnv) – The environment instance.

property active_functors: dict[str, list[str]]#

Name of active dataset functors by mode.

The keys are the modes and the values are the names of the dataset functors.

apply(mode, env_ids=None)[source]#

Apply dataset functors for the specified mode.

This method saves completed episodes by reading data from the environment’s episode buffers. It should be called before clearing the buffers during reset.

Parameters:
  • mode (str) – The mode to apply (currently only “save” is supported).

  • env_ids (Union[Sequence[int], Tensor, None]) – The indices of the environments to apply the functor to. Defaults to None, in which case the functor is applied to all environments.

Return type:

None

property available_modes: list[str]#

List of available modes for the dataset manager.

clear_cache()[source]#

Clear cached data from all dataset functors (for online training).

Iterates through all functors and clears their cache if they support online training mode (have clear_cache method).

Return type:

int

Returns:

Total number of cached items cleared across all functors.

close()[source]#

Finalize all dataset functors; repeated calls are safe.

Return type:

Optional[str]

property device: str#

Device on which to perform computations.

finalize()[source]#

Finalize every dataset functor exactly once.

Finalization is a storage barrier only; individual recorders must not implicitly commit live episode buffers here. All functors are attempted even when one fails, and their failures are reported together.

Return type:

Optional[str]

Returns:

Path to the first finalized dataset, or None if none was returned.

Raises:

RuntimeError – If one or more functors fail to finalize.

find_functors(name_keys)#

Find functors in the manager based on the names.

This function searches the manager for functors based on the names. The names can be specified as regular expressions or a list of regular expressions. The search is performed on the active functors in the manager.

Please check the resolve_matching_names() function for more information on the name matching.

Parameters:

name_keys (Union[str, Sequence[str]]) – A regular expression or a list of regular expressions to match the functor names.

Return type:

list[str]

Returns:

A list of functor names that match the input keys.

get_active_iterable_functors(env_idx)#

Returns the active functors as iterable sequence of tuples.

The first element of the tuple is the name of the functor and the second element is the raw value(s) of the functor.

Return type:

Sequence[tuple[str, Sequence[float]]]

Returns:

The active functors.

get_cached_data()[source]#

Get cached data from all dataset functors (for online training).

Iterates through all functors and collects cached data from those that support online training mode (have get_cached_data method).

Return type:

list[Dict[str, Any]]

Returns:

List of cached data dictionaries from all functors.

get_functor_cfg(functor_name)[source]#

Gets the configuration for the specified functor.

Parameters:

functor_name (str) – The name of the dataset functor.

Return type:

DatasetFunctorCfg

Returns:

The configuration of the dataset functor.

Raises:

ValueError – If the functor name is not found.

property num_envs: int#

Number of environments.

property requires_raw_actions: bool#

Whether a recorder needs pre-controller policy actions.

reset(env_ids=None)[source]#

Reset all dataset functors.

Parameters:

env_ids (Union[Sequence[int], Tensor, None]) – The environment ids. Defaults to None.

Return type:

dict[str, float]

Returns:

Empty dict (no logging info).

property save_failed_episodes: bool#

Whether any configured dataset recorder should keep failed episodes.

validate_policy_action(action)[source]#

Validate a live policy action for every descriptor-driven recorder.

Parameters:

action (Tensor) – Flat policy tensor before ActionManager processing.

Return type:

None

Action Terms#

class embodichain.lab.gym.envs.managers.ActionTerm[source]#

Bases: Functor

Base class for one independently processed and applied action slice.

contract_id identifies the stable action semantics; subclasses may change implementation without changing that public contract.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Resets the functor.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

command_type

Controller command or resource type written by the term.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
abstract property action_dim: int#

Number of raw policy-action dimensions owned by the term.

abstract property action_space: Box#

Raw policy-action bounds for this term.

abstract apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

abstract property command_type: str#

Controller command or resource type written by the term.

abstract property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

abstract property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

abstract process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

abstract property processed_actions: Tensor#

Current term-owned processed command rows.

abstract property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)#

Resets the functor.

Parameters:

env_ids (Optional[Sequence[int]]) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

resolved_contract_id()[source]#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

class embodichain.lab.gym.envs.managers.DefaultJointPositionAction[source]#

Bases: _JointAction

Map normalized actions around a configured default joint pose.

clip defaults to 1.0; None preserves the unbounded behavior of legacy DefaultJointPositionTerm snapshots.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Clear selected action history and commands without changing bias.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

previous_raw_actions

Return raw actions from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property previous_raw_actions: Tensor#

Return raw actions from the previous control step.

process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)[source]#

Clear selected action history and commands without changing bias.

Parameters:

env_ids (list[int] | Tensor | None) – Rows to reset. None selects all rows.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

class embodichain.lab.gym.envs.managers.EefPoseAction[source]#

Bases: ActionTerm

Convert an absolute EEF pose into selected-arm joint targets.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Resets the functor.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

ik_success

Return per-environment IK success for the latest request.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

part_name

Return the robot control part used for IK and FK.

previous_raw_actions

Return pose requests from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property ik_success: Tensor#

Return per-environment IK success for the latest request.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property part_name: str#

Return the robot control part used for IK and FK.

property previous_raw_actions: Tensor#

Return pose requests from the previous control step.

process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)[source]#

Resets the functor.

Parameters:

env_ids (list[int] | Tensor | None) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

RelativeEefPoseAction defaults to the configured arm IK solver. Both backends validate the solver chain’s joint order and map seeds and returned commands between that order and the configured joint selection. Its optional ik_method="dls" backend applies one damped Jacobian increment in the arena frame, rescales the whole increment to max_joint_delta, and clips position limits, reading current robot limits on every control step. It uses the actual solver root-link transform, and shifts the end-link Jacobian to the configured TCP offset on the environment device. This backend approximates a local pose step; ik_success means a finite controller step and does not assert pose convergence. The action descriptor records its method, damping, bound, and success semantics. joint_delta, joint_step_limited, and joint_limit_clipped_mask expose the applied increment and both forms of limiting; resets clear selected rows.

class embodichain.lab.gym.envs.managers.RelativeEefPoseAction[source]#

Bases: EefPoseAction

Convert normalized arena-frame pose deltas into selected-arm targets.

Each request starts from FK of the measured joint positions. Translation adds translation_scale * clamp(u[:3], -1, 1) in meters; rotation left-multiplies the measured rotation by the exponential map of rotation_scale * clamp(u[3:], -1, 1) in radians. There is no accumulated pose target. Failed or nonfinite IK solutions hold the measured arm joints. Optional ik_method="dls" applies one damped, bounded local Jacobian increment. Its success flag means a finite controller step, not convergence to the requested pose.

Methods:

__init__(cfg, env)

Initialize the measured-pose delta controller.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Resolve one finite normalized request from current measured joints.

reset([env_ids])

Clear selected action history and diagnostics.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

clipped_mask

Return which components exceeded normalized bounds on the latest step.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

current_pose

Return latest measured arena-frame transforms, or identity after reset.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

ik_success

Return per-environment IK success for the latest request.

joint_delta

Return the latest commanded increment in selected joint order.

joint_limit_clipped_mask

Return DLS joint components clipped to position limits.

joint_step_limited

Return DLS rows rescaled to respect the maximum joint increment.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

part_name

Return the robot control part used for IK and FK.

previous_raw_actions

Return pose requests from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

target_pose

Return latest requested arena-frame transforms, or identity after reset.

__init__(cfg, env)[source]#

Initialize the measured-pose delta controller.

Parameters:
  • cfg (ActionTermCfg) – Term config with part_name, optional positive finite translation_scale (default 0.02 m), and rotation_scale (default 0.10 rad), applied per component per control step. ik_method selects the existing "solver" backend or one "dls" step with positive finite dls_damping (0.05) and max_joint_delta (0.10, in each joint’s position units).

  • env (EmbodiedEnv) – Environment supplying measured joints, FK, and IK.

property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()#

Apply the latest processed command to the owned resource.

Return type:

None

property clipped_mask: Tensor#

Return which components exceeded normalized bounds on the latest step.

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property current_pose: Tensor#

Return latest measured arena-frame transforms, or identity after reset.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property ik_success: Tensor#

Return per-environment IK success for the latest request.

property joint_delta: Tensor#

Return the latest commanded increment in selected joint order.

property joint_limit_clipped_mask: Tensor#

Return DLS joint components clipped to position limits.

property joint_step_limited: Tensor#

Return DLS rows rescaled to respect the maximum joint increment.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property part_name: str#

Return the robot control part used for IK and FK.

property previous_raw_actions: Tensor#

Return pose requests from the previous control step.

process_actions(actions)[source]#

Resolve one finite normalized request from current measured joints.

Parameters:

actions (Tensor) – Floating requests with shape (num_envs, 6).

Raises:
  • TypeError – If requests are not floating tensors.

  • ValueError – If their shape/device is invalid or any value is nonfinite.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)[source]#

Clear selected action history and diagnostics.

Parameters:

env_ids (list[int] | Tensor | None) – Environment rows to reset; None selects every row.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

property target_pose: Tensor#

Return latest requested arena-frame transforms, or identity after reset.

class embodichain.lab.gym.envs.managers.JointEffortAction[source]#

Bases: _JointAction

Apply scaled effort actions to selected joints.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Resets the functor.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

previous_raw_actions

Return raw actions from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property previous_raw_actions: Tensor#

Return raw actions from the previous control step.

process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)#

Resets the functor.

Parameters:

env_ids (list[int] | Tensor | None) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

class embodichain.lab.gym.envs.managers.JointPositionAction[source]#

Bases: _JointAction

Apply absolute position actions to selected joints.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Resets the functor.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

previous_raw_actions

Return raw actions from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property previous_raw_actions: Tensor#

Return raw actions from the previous control step.

process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)#

Resets the functor.

Parameters:

env_ids (list[int] | Tensor | None) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

class embodichain.lab.gym.envs.managers.JointPositionToLimitsAction[source]#

Bases: _JointAction

Map normalized selected-joint actions to position limits.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Resets the functor.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

previous_raw_actions

Return raw actions from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property previous_raw_actions: Tensor#

Return raw actions from the previous control step.

process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)#

Resets the functor.

Parameters:

env_ids (list[int] | Tensor | None) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

class embodichain.lab.gym.envs.managers.JointVelocityAction[source]#

Bases: _JointAction

Apply scaled velocity actions to selected joints.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Resets the functor.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

previous_raw_actions

Return raw actions from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property previous_raw_actions: Tensor#

Return raw actions from the previous control step.

process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)#

Resets the functor.

Parameters:

env_ids (list[int] | Tensor | None) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

class embodichain.lab.gym.envs.managers.ParallelGripperAction[source]#

Bases: ActionTerm

Map one scalar policy action to explicit parallel-gripper joints.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Resets the functor.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

close_command

Return the resolved close joint command.

command_mode

Return the continuous or binary scalar command mode.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

lower_command

Return the continuous command at normalized value -1.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

open_command

Return the resolved open joint command.

part_name

Return the robot control part owning the gripper joints.

previous_raw_actions

Return scalar requests from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

upper_command

Return the continuous command at normalized value 1.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

property close_command: Tensor#

Return the resolved close joint command.

property command_mode: str#

Return the continuous or binary scalar command mode.

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property lower_command: Tensor#

Return the continuous command at normalized value -1.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property open_command: Tensor#

Return the resolved open joint command.

property part_name: str#

Return the robot control part owning the gripper joints.

property previous_raw_actions: Tensor#

Return scalar requests from the previous control step.

process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)[source]#

Resets the functor.

Parameters:

env_ids (list[int] | Tensor | None) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

property upper_command: Tensor#

Return the continuous command at normalized value 1.

class embodichain.lab.gym.envs.managers.RelativeJointPositionAction[source]#

Bases: _JointAction

Add scaled relative actions to current selected positions.

Methods:

__init__(cfg, env)

Initialize an action term.

apply_actions()

Apply the latest processed command to the owned resource.

process_actions(actions)

Process one manager-owned raw action slice.

reset([env_ids])

Resets the functor.

resolved_contract_id()

Return the configured contract or the implementation's stable ID.

serialize()

General serialization call.

Attributes:

action_dim

Number of raw policy-action dimensions owned by the term.

action_space

Raw policy-action bounds for this term.

controlled_joint_ids

Ordered robot joint IDs owned by the term.

descriptor

Describe the raw policy-action representation.

device

Device on which to perform computations.

name

Return the manager-bound, read-only configuration name.

num_envs

Number of environments.

previous_raw_actions

Return raw actions from the previous control step.

processed_actions

Current term-owned processed command rows.

raw_actions

Current raw policy-action rows.

__init__(cfg, env)[source]#

Initialize an action term.

Parameters:
property action_dim: int#

Number of raw policy-action dimensions owned by the term.

property action_space: Box#

Raw policy-action bounds for this term.

apply_actions()[source]#

Apply the latest processed command to the owned resource.

Return type:

None

property controlled_joint_ids: tuple[int, ...]#

Ordered robot joint IDs owned by the term.

property descriptor: ActionTermDescriptor#

Describe the raw policy-action representation.

property device: str#

Device on which to perform computations.

property name: str#

Return the manager-bound, read-only configuration name.

property num_envs: int#

Number of environments.

property previous_raw_actions: Tensor#

Return raw actions from the previous control step.

process_actions(actions)[source]#

Process one manager-owned raw action slice.

Return type:

None

property processed_actions: Tensor#

Current term-owned processed command rows.

property raw_actions: Tensor#

Current raw policy-action rows.

reset(env_ids=None)#

Resets the functor.

Parameters:

env_ids (list[int] | Tensor | None) – The environment ids. Defaults to None, in which case all environments are considered.

Return type:

None

resolved_contract_id()#

Return the configured contract or the implementation’s stable ID.

Raises:

ValueError – If a configured contract differs from the implementation contract declared by the action class.

Return type:

str | None

serialize()#

General serialization call. Includes the configuration dict.

Return type:

dict

Action Descriptors#

class embodichain.lab.gym.envs.managers.ActionDescriptor[source]#

Bind one action-term descriptor to a manager-owned flat slice.

Parameters:
  • name (str) – Stable term name from the ActionManager configuration.

  • start (int) – Inclusive flat-action offset.

  • stop (int) – Exclusive flat-action offset.

  • term (ActionTermDescriptor) – Term-local action semantics.

Methods:

__init__(name, start, stop, term)

to_dict()

Return an owned JSON-compatible representation.

__init__(name, start, stop, term)#
to_dict()[source]#

Return an owned JSON-compatible representation.

Return type:

dict[str, None | bool | int | float | str | list[None | bool | int | float | str | list[JSONValue] | dict[str, JSONValue]] | dict[str, None | bool | int | float | str | list[JSONValue] | dict[str, JSONValue]]]

class embodichain.lab.gym.envs.managers.ActionTermDescriptor[source]#

Describe one term’s raw policy-action representation.

Parameters:
  • representation (str) – Stable semantic representation identifier.

  • action_dim (int) – Number of flat raw policy-action dimensions.

  • feature_names (tuple[str, ...]) – Ordered name for each raw dimension.

  • units (tuple[str, ...]) – Ordered unit label for each raw dimension.

  • normalization (str | None) – Optional normalization identifier.

  • joint_names (tuple[str, ...]) – Ordered joints controlled by the term.

  • metadata (Mapping[str, None | bool | int | float | str | list[None | bool | int | float | str | list[JSONValue] | dict[str, JSONValue]] | dict[str, None | bool | int | float | str | list[JSONValue] | dict[str, JSONValue]]]) – Additional JSON-compatible representation metadata.

  • contract (str | None) – Versioned, implementation-independent action contract ID.

Methods:

__init__(representation, action_dim, ...[, ...])

__new__(**kwargs)

__init__(representation, action_dim, feature_names, units, normalization, joint_names, metadata=<factory>, contract=None)#
__new__(**kwargs)#
class embodichain.lab.gym.envs.managers.ActionTrace[source]#

Capture one policy request and its controller-side action stages.

requested is the flat action emitted by a policy. processed and executed are term-keyed controller commands. The latter is optional because a trace can be captured before the controller boundary is applied. Keeping these stages explicit prevents dataset consumers from assigning different meanings to one overloaded actions buffer.

Methods:

__init__(requested, processed, executed, ...)

clone()

Return an independently owned copy of this trace.

__init__(requested, processed, executed, command_types, controlled_joint_ids)#
clone()[source]#

Return an independently owned copy of this trace.

Return type:

ActionTrace

Observation Functions#

Classes:

compute_exteroception

Compute the exteroception for the observation space.

get_articulation_joint_drive

Get the joint drive properties of the articulation in the environment with caching.

get_rigid_object_physics_attributes

Get the physics attributes of the rigid object in the environment with caching.

Functions:

compute_semantic_mask(env, obs, entity_cfg, ...)

Compute the semantic mask for the specified scene entity.

get_object_body_scale(env, obs, entity_cfg)

Get the body scale of the objects in the environment.

get_object_pose(env, obs, entity_cfg[, ...])

Get the arena poses of the objects in the environment.

get_object_uid(env, obs, entity_cfg)

Get the user IDs of the objects in the environment.

get_rigid_object_pose(env, obs, entity_cfg)

Get the world poses of the rigid objects in the environment.

get_rigid_object_velocity(env, obs, entity_cfg)

Get the world velocities of the rigid objects in the environment.

get_robot_eef_pose(env, obs[, part_name, ...])

Get robot end-effector pose using forward kinematics.

get_sensor_intrinsics(env, obs, entity_cfg)

Get the intrinsic matrix of a sensor (camera).

get_sensor_pose_in_robot_frame(env, obs, ...)

Get the pose of a sensor in the robot's base coordinate frame.

normalize_robot_joint_data(env, data[, ...])

Normalize selected observation columns using their robot joint limits.

target_position(env, obs[, target_pose_key])

Get virtual target position from env state.

class embodichain.lab.gym.envs.managers.observations.compute_exteroception[source]#

Compute the exteroception for the observation space.

The exteroception is currently defined as a set of keypoints around a reference pose, which are prjected from 3D space to 2D image plane. The reference pose can derive from the following sources:

  • Pose from robot control part (e.g., end-effector, usually tcp pose)

  • Object affordance pose (e.g., handle pose of a mug or a pick pose of a cube)

Therefore, the exteroception are defined in the camera-like sensor, for example. descriptor = {

“cam_high”: [
{

“type”: “affordance”, “obj_uid”: “obj1”, “key”: “grasp_pose”, “is_arena_coord”: True

}, {

“type”: “affordance”, “obj_uid”: “obj1”, “key”: “place_pose”,

}, {

“type”: “robot”, “control_part”: “left_arm”,

}, {

“type”: “robot”, “control_part”: “right_arm”,

}

}

Explanation of the parameters:
  • The key of the dictionary is the sensor uid.

  • The value is another dictionary, where the key is the source type, and the value is a dictionary of parameters.

  • For affordance source type, the parameters are:
    • obj_uid: The uid of the object to get the affordance pose from.

    • key: The key of the affordance pose in the affordance data.

    • is_arena_coord: Whether the affordance pose is in the arena coordinate system. Default is False.

  • For robot source type, the parameters are:
    • control_part: The control part of the robot to get the pose from.

Methods:

__init__(cfg, env)

Initialize the functor.

expand_pose(pose, x_interval, y_interval, ...)

Expand pose with keypoints along x and y axes.

shift_pose(pose, axis, shift)

Shift the pose along the specified axis by the given amount.

__init__(cfg, env)[source]#

Initialize the functor.

Parameters:
static expand_pose(pose, x_interval, y_interval, kpnts_number, ref_pose=None)[source]#

Expand pose with keypoints along x and y axes.

Parameters:
  • pose (Tensor) – The original pose tensor of shape (B, 4, 4).

  • x_interval (float) – The interval for expanding along x-axis.

  • y_interval (float) – The interval for expanding along y-axis.

  • kpnts_number (int) – Number of keypoints to generate for each axis.

  • ref_pose (Tensor) – Reference pose tensor of shape (B, 4, 4). If None, uses identity matrix.

Return type:

Tensor

Returns:

Expanded poses tensor of shape (B, 1 + 2*kpnts_number, 4, 4).

static shift_pose(pose, axis, shift)[source]#

Shift the pose along the specified axis by the given amount.

Parameters:
  • pose (Tensor) – The original pose tensor of shape (B, 4, 4).

  • axis (int) – The axis along which to shift (0 for x, 1 for y, 2 for z).

  • shift (float) – The amount to shift along the specified axis.

Return type:

Tensor

embodichain.lab.gym.envs.managers.observations.compute_semantic_mask(env, obs, entity_cfg, foreground_uids, is_right=False)[source]#

Compute the semantic mask for the specified scene entity.

Note

The semantic mask is defined as (B, H, W, len(SemanticMask)) where these channels represents: - background channel: the instance id of the background is set to 1 (0 if not background) - foreground channel: the instance id of the foreground objects is set to 1 (0 if not foreground) - robot left-side channel: the instance id of the robot left-side is set to 1 - robot right-side channel: the instance id of the robot right-side is set to 1

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity.

  • foreground_uids (Sequence[str]) – The list of uids for the foreground objects.

  • is_right (bool) – Whether to use the right camera for stereo cameras. Default is False. Only applicable if the sensor is a StereoCamera.

Return type:

Tensor

Returns:

A tensor of shape (num_envs, height, width) representing the semantic mask.

class embodichain.lab.gym.envs.managers.observations.get_articulation_joint_drive[source]#

Get the joint drive properties of the articulation in the environment with caching.

This functor retrieves and caches joint drive properties (stiffness, damping, max_effort, max_velocity, friction) for articulations (including robots). The cache is cleared when the environment resets, ensuring fresh values are fetched at the start of each episode.

If the articulation with the specified UID does not exist in the environment, a zero tensor will be returned for each attribute.

The cached data is stored per entity UID. When called, if data is cached, it returns a clone of the cached tensor to prevent accidental modifications.

Note

Joint drive properties are typically constant during an episode, so caching improves performance by avoiding repeated queries.

Parameters:

Methods:

__init__(cfg, env)

Initialize the joint drive functor.

reset([env_ids])

Clear the cached joint drive properties.

__init__(cfg, env)[source]#

Initialize the joint drive functor.

Parameters:
reset(env_ids=None)[source]#

Clear the cached joint drive properties.

Parameters:

env_ids (Optional[Sequence[int]]) – The environment ids. Defaults to None, which clears all cache.

Return type:

None

embodichain.lab.gym.envs.managers.observations.get_object_body_scale(env, obs, entity_cfg)[source]#

Get the body scale of the objects in the environment.

If the object with the specified UID does not exist in the environment, a zero tensor will be returned.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity.

Return type:

Tensor

Returns:

A tensor of shape (num_envs, 3) representing the body scale of the objects.

embodichain.lab.gym.envs.managers.observations.get_object_pose(env, obs, entity_cfg, to_matrix=True)[source]#

Get the arena poses of the objects in the environment.

If the object with the specified UID does not exist in the environment, a zero tensor will be returned.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity.

  • to_matrix (bool) – Whether to return the pose as a 4x4 transformation matrix. If False, returns (x, y, z, qx, qy, qz, qw).

Return type:

Tensor

Returns:

A tensor of shape (num_envs, 7) or (num_envs, 4, 4) representing the world poses of the objects.

embodichain.lab.gym.envs.managers.observations.get_object_uid(env, obs, entity_cfg)[source]#

Get the user IDs of the objects in the environment.

If the object with the specified UID does not exist in the environment, a zero tensor will be returned.

Note

  • If asset is RigidObject, the user IDs is shaped as (num_envs,)

  • If asset is Articulation or Robot, the user IDs is shaped as (num_envs, num_links) and ordered by

    link_names in the configuration.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity.

Return type:

Tensor

Returns:

A tensor of shape (num_envs,) representing the user IDs of the objects.

class embodichain.lab.gym.envs.managers.observations.get_rigid_object_physics_attributes[source]#

Get the physics attributes of the rigid object in the environment with caching.

This functor retrieves and caches physics attributes (mass, friction, damping, inertia) for rigid objects. The cache is cleared when the environment resets, ensuring fresh values are fetched at the start of each episode.

If the rigid object with the specified UID does not exist in the environment, a zero tensor will be returned for each attribute.

The cached data is stored per entity UID. When called, if data is cached, it returns a clone of the cached tensor to prevent accidental modifications.

Note

Physics attributes are typically constant during an episode, so caching improves performance by avoiding repeated queries to the physics engine.

Parameters:

Methods:

__init__(cfg, env)

Initialize the physics attributes functor.

reset([env_ids])

Clear the cached physics attributes.

__init__(cfg, env)[source]#

Initialize the physics attributes functor.

Parameters:
reset(env_ids=None)[source]#

Clear the cached physics attributes.

Parameters:

env_ids (Optional[Sequence[int]]) – The environment ids. Defaults to None, which clears all cache.

Return type:

None

embodichain.lab.gym.envs.managers.observations.get_rigid_object_pose(env, obs, entity_cfg, to_matrix=True)[source]#

Get the world poses of the rigid objects in the environment.

If the rigid object with the specified UID does not exist in the environment, a zero tensor will be returned.

Note

This method will be deprecated in the future and replaced by get_object_pose as the distinction between rigid objects and general objects is being removed. Please use get_object_pose instead when possible.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity.

  • to_matrix (bool) – Whether to return the pose as a 4x4 transformation matrix. If False, returns (x, y, z, qx, qy, qz, qw).

Return type:

Tensor

Returns:

A tensor of shape (num_envs, 7) or (num_envs, 4, 4) representing the world poses of the rigid objects.

embodichain.lab.gym.envs.managers.observations.get_rigid_object_velocity(env, obs, entity_cfg)[source]#

Get the world velocities of the rigid objects in the environment.

If the rigid object with the specified UID does not exist in the environment, a zero tensor will be returned.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity.

Return type:

Tensor

Returns:

A tensor of shape (num_envs, 6) representing the linear and angular velocities of the rigid objects.

embodichain.lab.gym.envs.managers.observations.get_robot_eef_pose(env, obs, part_name=None, position_only=False)[source]#

Get robot end-effector pose using forward kinematics.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • part_name (str | None) – The name of the control part. If None, uses default part.

  • position_only (bool) – If True, returns only position (3D). If False, returns full pose (4x4 matrix).

Return type:

Tensor

Returns:

A tensor of shape (num_envs, 3) if position_only=True, or (num_envs, 4, 4) otherwise.

embodichain.lab.gym.envs.managers.observations.get_sensor_intrinsics(env, obs, entity_cfg, is_right=False)[source]#

Get the intrinsic matrix of a sensor (camera).

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • entity_cfg (SceneEntityCfg) – The configuration of the sensor entity.

  • is_right (bool) – Whether to return the right camera intrinsics for stereo cameras. Defaults to False (left camera). Ignored for monocular cameras.

Return type:

Tensor

Returns:

A tensor of shape (num_envs, 3, 3) representing the camera intrinsics.

embodichain.lab.gym.envs.managers.observations.get_sensor_pose_in_robot_frame(env, obs, entity_cfg, robot_uid=None, is_right=False)[source]#

Get the pose of a sensor in the robot’s base coordinate frame.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – The observation dictionary.

  • entity_cfg (SceneEntityCfg) – The configuration of the sensor entity.

  • robot_uid (str | None) – The uid of the robot. If None, uses the default robot from env.

  • is_right (bool) – Whether to return the right camera intrinsics for stereo cameras. Defaults to False (left camera). Ignored for monocular cameras.

Return type:

Tensor

Returns:

A tensor of shape (num_envs, 4, 4) representing the sensor pose in robot coordinates as a transformation matrix.

embodichain.lab.gym.envs.managers.observations.normalize_robot_joint_data(env, data, joint_ids=None, limit='qpos_limits', range=[0.0, 1.0], *, control_parts=None)[source]#

Normalize selected observation columns using their robot joint limits.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • data (Tensor) – Joint data in the environment’s active-joint observation order.

  • joint_ids (Optional[Sequence[int]]) – Observation-column indices to normalize. Mutually exclusive with control_parts.

  • limit (Literal['qpos_limits', 'qvel_limits']) – The type of joint limits to be used for normalization. Options are: - qpos_limits: Use the joint position limits for normalization. - qvel_limits: Use the joint velocity limits for normalization.

  • range (Sequence[float]) – The range to normalize the joint data to. Defaults to [0.0, 1.0].

  • control_parts (Optional[Sequence[str]]) – Robot control parts whose observed joints are normalized. Joint names are resolved against the current robot and mapped into env.active_joint_ids, including observed mimic joints.

Return type:

Tensor

Returns:

The input tensor with the selected columns normalized in place.

Raises:

ValueError – If neither or both selectors are given, or the selected control parts have no joints in the observation.

embodichain.lab.gym.envs.managers.observations.target_position(env, obs, target_pose_key='goal_pose')[source]#

Get virtual target position from env state.

Reads target pose from env.{target_pose_key} (set by randomize_target_pose event). Returns zeros if not yet initialized (e.g., during env initialization before reset).

Parameters:
  • env (EmbodiedEnv) – The environment instance

  • obs (TensorDict[str, Union[Tensor, TensorDict[str, Tensor]]]) – Observation dict (unused, for API compatibility)

  • target_pose_key (str) – Key for target pose in env (default: “goal_pose”)

Return type:

Tensor

Returns:

Target position tensor of shape (num_envs, 3). Returns zeros if target_pose_key is not found (e.g., before first reset).

Reward Functions#

Common reward functors for reinforcement learning tasks.

Functions:

action_smoothness_penalty(env, obs, action, info)

Penalize large action changes between consecutive timesteps.

distance_between_objects(env, obs, action, info)

Reward based on distance between two rigid objects.

distance_to_target(env, obs, action, info[, ...])

Reward based on absolute distance to a virtual target pose.

incremental_distance_to_target(env, obs, ...)

Incremental reward for progress toward a virtual target pose.

joint_limit_penalty(env, obs, action, info)

Penalize robot joints that are close to their position limits.

joint_velocity_penalty(env, obs, action, info)

Penalize high joint velocities to encourage smooth motion.

orientation_alignment(env, obs, action, info)

Reward rotational alignment between two rigid objects.

reaching_behind_object(env, obs, action, info)

Reward for positioning end-effector behind object for pushing.

success_reward(env, obs, action, info)

Sparse bonus reward when task succeeds.

embodichain.lab.gym.envs.managers.rewards.action_smoothness_penalty(env, obs, action, info)[source]#

Penalize large action changes between consecutive timesteps.

Encourages smooth control commands by penalizing sudden changes in actions. Reads the previous action from the RL rollout_buffer (action / done). Returns zeros when that buffer is unavailable (e.g. evaluation).

Return type:

Tensor

Returns:

Penalty tensor of shape (num_envs,). Zero on first step (no previous action), negative on subsequent steps (larger change = more negative).

Example

```json {

“func”: “action_smoothness_penalty”, “weight”: 0.01, “params”: {}

}#

embodichain.lab.gym.envs.managers.rewards.distance_between_objects(env, obs, action, info, source_entity_cfg=None, target_entity_cfg=None, exponential=False, sigma=1.0)[source]#

Reward based on distance between two rigid objects.

Encourages the source object to get closer to the target object. Can use either linear negative distance or exponential Gaussian-shaped reward.

Parameters:
  • source_entity_cfg (SceneEntityCfg) – Configuration for the source object (e.g., {“uid”: “cube”})

  • target_entity_cfg (SceneEntityCfg) – Configuration for the target object (e.g., {“uid”: “goal_sphere”})

  • exponential (bool) – If True, use exponential reward exp(-d²/2σ²), else use -distance

  • sigma (float) – Standard deviation for exponential reward (controls reward spread)

Return type:

Tensor

Returns:

Reward tensor of shape (num_envs,). Higher when objects are closer. - Linear mode: ranges from -inf to 0 (0 when objects touch) - Exponential mode: ranges from 0 to 1 (1 when objects touch)

Example

```json {

“func”: “distance_between_objects”, “weight”: 0.5, “params”: {

“source_entity_cfg”: {“uid”: “cube”}, “target_entity_cfg”: {“uid”: “target”}, “exponential”: true, “sigma”: 0.2

}

}#

embodichain.lab.gym.envs.managers.rewards.distance_to_target(env, obs, action, info, source_entity_cfg=None, target_pose_key='target_pose', exponential=False, sigma=1.0, use_xy_only=False)[source]#

Reward based on absolute distance to a virtual target pose.

Encourages an object to get closer to a target pose specified in the info dict. Unlike incremental_distance_to_target, this provides direct distance-based reward.

Parameters:
  • source_entity_cfg (SceneEntityCfg) – Configuration for the object (e.g., {“uid”: “cube”})

  • target_pose_key (str) – Key in info dict for target pose (default: “target_pose”) Can be (num_envs, 3) position or (num_envs, 4, 4) transform

  • exponential (bool) – If True, use exponential reward exp(-d²/2σ²), else use -distance

  • sigma (float) – Standard deviation for exponential reward (default: 1.0)

  • use_xy_only (bool) – If True, ignore z-axis and only consider horizontal distance

Return type:

Tensor

Returns:

Reward tensor of shape (num_envs,). - Linear mode: -distance (negative, approaches 0 when close) - Exponential mode: exp(-d²/2σ²) (0 to 1, approaches 1 when close)

Example

```json {

“func”: “distance_to_target”, “weight”: 0.5, “params”: {

“source_entity_cfg”: {“uid”: “cube”}, “target_pose_key”: “goal_pose”, “exponential”: false, “use_xy_only”: true

}

}#

embodichain.lab.gym.envs.managers.rewards.incremental_distance_to_target(env, obs, action, info, source_entity_cfg=None, target_pose_key='target_pose', tanh_scale=10.0, positive_weight=1.0, negative_weight=1.0, use_xy_only=False)[source]#

Incremental reward for progress toward a virtual target pose.

Rewards the robot for getting closer to the target compared to previous timestep. Stores previous distance in env._reward_states for comparison. Uses tanh shaping to normalize rewards and supports asymmetric weighting for approach vs. retreat.

Parameters:
  • source_entity_cfg (SceneEntityCfg) – Configuration for the object (e.g., {“uid”: “cube”})

  • target_pose_key (str) – Key for target pose in env (default: “target_pose”) Reads from env._{target_pose_key} set by randomize_target_pose event Can be (num_envs, 3) position or (num_envs, 4, 4) transform

  • tanh_scale (float) – Scaling for tanh normalization (higher = more sensitive, default: 10.0)

  • positive_weight (float) – Multiplier for reward when getting closer (default: 1.0)

  • negative_weight (float) – Multiplier for penalty when moving away (default: 1.0)

  • use_xy_only (bool) – If True, ignore z-axis and only consider horizontal distance

Returns:

  • Positive when getting closer (scaled by positive_weight)

  • Negative when moving away (scaled by negative_weight)

  • Magnitude bounded by tanh function

Return type:

Reward tensor of shape (num_envs,). Zero on first call, then

Note

This function maintains state using env._reward_states[f”prev_dist_{uid}_{key}”]. State is automatically reset when the environment resets.

Example

```json {

“func”: “incremental_distance_to_target”, “weight”: 1.0, “params”: {

“source_entity_cfg”: {“uid”: “cube”}, “target_pose_key”: “goal_pose”, “tanh_scale”: 10.0, “positive_weight”: 2.0, “negative_weight”: 0.5, “use_xy_only”: true

}

}#

embodichain.lab.gym.envs.managers.rewards.joint_limit_penalty(env, obs, action, info, robot_uid='robot', joint_ids=slice(None, None, None), margin=0.1)[source]#

Penalize robot joints that are close to their position limits.

Prevents joints from reaching their physical limits, which can cause instability or singularities. Penalty increases as joints approach limits within the margin.

Parameters:
  • robot_uid (str) – Robot entity UID in simulation (default: “robot”)

  • joint_ids (slice | list[int]) – Joint indices to monitor (default: all joints)

  • margin (float) – Normalized distance threshold (0 to 1). Penalty applied when joint is within this fraction of its range from either limit. Example: 0.1 means penalty when within 10% of limits.

Return type:

Tensor

Returns:

Penalty tensor of shape (num_envs,). Always negative or zero. Sum of penalties across all monitored joints.

Example

```json {

“func”: “joint_limit_penalty”, “weight”: 0.01, “params”: {

“robot_uid”: “robot”, “joint_ids”: [0, 1, 2, 3, 4, 5], “margin”: 0.1

}

}#

embodichain.lab.gym.envs.managers.rewards.joint_velocity_penalty(env, obs, action, info, robot_uid='robot', joint_ids=None, part_name=None)[source]#

Penalize high joint velocities to encourage smooth motion.

Computes the L2 norm of joint velocities and returns negative value as penalty. Useful for preventing jerky or unstable robot movements.

Parameters:
  • robot_uid (str) – Robot entity UID in simulation (default: “robot”)

  • joint_ids (slice | list[int] | None) – Specific joint indices to penalize. Takes priority over part_name. Example: [0, 1, 2] or slice(0, 6)

  • part_name (str | None) – Control part name (e.g., “arm”). Used only if joint_ids is None. Will penalize all joints in the specified part.

Return type:

Tensor

Returns:

Penalty tensor of shape (num_envs,). Always negative or zero. Magnitude increases with joint velocity (larger velocity = more negative).

Example

```json {

“func”: “joint_velocity_penalty”, “weight”: 0.001, “params”: {

“robot_uid”: “robot”, “part_name”: “arm”

}

}#

embodichain.lab.gym.envs.managers.rewards.orientation_alignment(env, obs, action, info, source_entity_cfg=None, target_entity_cfg=None)[source]#

Reward rotational alignment between two rigid objects.

Encourages the source object’s orientation to match the target object’s orientation. Uses rotation matrix trace to measure alignment.

Parameters:
  • source_entity_cfg (SceneEntityCfg) – Configuration for the source object (e.g., {“uid”: “cube”})

  • target_entity_cfg (SceneEntityCfg) – Configuration for the target object (e.g., {“uid”: “reference”})

Return type:

Tensor

Returns:

Reward tensor of shape (num_envs,). Ranges from -1 to 1. - 1.0: Perfect alignment (same orientation) - 0.0: 90° rotation difference - -1.0: 180° rotation difference (opposite orientation)

Example

```json {

“func”: “orientation_alignment”, “weight”: 0.5, “params”: {

“source_entity_cfg”: {“uid”: “object”}, “target_entity_cfg”: {“uid”: “goal_object”}

}

}#

embodichain.lab.gym.envs.managers.rewards.reaching_behind_object(env, obs, action, info, object_cfg=None, target_pose_key='goal_pose', behind_offset=0.015, height_offset=0.015, distance_scale=5.0, part_name=None)[source]#

Reward for positioning end-effector behind object for pushing.

Encourages the robot’s end-effector to reach a position behind the object along the object-to-goal direction. Useful for push manipulation tasks.

Parameters:
  • object_cfg (SceneEntityCfg) – Configuration for the object to push (e.g., {“uid”: “cube”})

  • target_pose_key (str) – Key in info dict for goal pose (default: “goal_pose”) Can be (num_envs, 3) position or (num_envs, 4, 4) transform

  • behind_offset (float) – Distance behind object to reach (in meters, default: 0.015)

  • height_offset (float) – Additional height above object (in meters, default: 0.015)

  • distance_scale (float) – Scaling factor for tanh function (higher = steeper, default: 5.0)

  • part_name (str) – Robot part name for FK computation (e.g., “arm”)

Return type:

Tensor

Returns:

Reward tensor of shape (num_envs,). Ranges from 0 to 1. - 1.0: End-effector at ideal pushing position - 0.0: End-effector far from ideal position

Example

```json {

“func”: “reaching_behind_object”, “weight”: 0.1, “params”: {

“object_cfg”: {“uid”: “cube”}, “target_pose_key”: “goal_pose”, “behind_offset”: 0.015, “height_offset”: 0.015, “distance_scale”: 5.0, “part_name”: “arm”

}

}#

embodichain.lab.gym.envs.managers.rewards.success_reward(env, obs, action, info)[source]#

Sparse bonus reward when task succeeds.

Provides a fixed reward when the task success condition is met. Reads success status from info[‘success’] which should be set by the environment.

Return type:

Tensor

Returns:

Reward tensor of shape (num_envs,). - 1.0 when successful - 0.0 when not successful or if ‘success’ key missing

Note

The environment’s get_info() must populate info[‘success’] with a boolean tensor indicating success status for each environment.

Example

```json {

“func”: “success_reward”, “weight”: 10.0, “params”: {}

}#

Event Functions#

Functions:

apply_expansion_profile_reset(env, env_ids, *)

Apply one expansion candidate's scene and visual state at reset.

create_rigid_constraint(env, env_ids, ...[, ...])

Attach two rigid objects via a fixed constraint for the given env_ids.

drop_rigid_object_group_sequentially(env, ...)

Drop rigid object group from a specified height sequentially in the environment.

register_entity_attrs(env, env_ids, entity_cfg)

Register the atrributes of an entity to the env.registration dict.

register_info_to_env(env, env_ids, registry)

Register entity information without advancing physics by default.

remove_rigid_constraint(env, env_ids, name)

Remove the named constraint for the given env_ids.

set_detached_uids_for_env_reset(env, ...)

Set the UIDs of objects that are detached from automatic reset in the environment.

wait_for_dynamic_objects_to_settle(env, env_ids)

Advance physics until selected dynamic objects remain stationary.

Classes:

prepare_extra_attr

replace_assets_from_group

Replace a directory-sampled mesh or a prepared semantic asset combination.

embodichain.lab.gym.envs.managers.events.apply_expansion_profile_reset(env, env_ids, *, cube_pose=None, visual_registry=None, visual_assignments=None, visual_seed=0)[source]#

Apply one expansion candidate’s scene and visual state at reset.

The expansion runner owns candidate enumeration and supplies the resolved batch payload through the event configuration immediately before reset. This event owns only the environment mutation boundary, so the same reset lifecycle is used for ordinary and expansion environments.

Parameters:
  • env (EmbodiedEnv) – Environment whose selected rows should be prepared.

  • env_ids (Tensor | Sequence[int] | slice | None) – Environment rows selected by the reset lifecycle.

  • cube_pose (Tensor | None) – Optional (num_envs, 7) pose tensor in position/xyzw form.

  • visual_registry (Any | None) – Optional registry exposing apply_to_environment.

  • visual_assignments (Mapping[int, str] | None) – Row-to-profile assignments for the visual registry.

  • visual_seed (int) – Base seed passed to the visual registry.

Return type:

None

embodichain.lab.gym.envs.managers.events.create_rigid_constraint(env, env_ids, obj_a_cfg, obj_b_cfg, name, local_frame_a=None, local_frame_b=None)[source]#

Attach two rigid objects via a fixed constraint for the given env_ids.

Registered under a custom event mode (e.g. "attach"); the task triggers it with env.event_manager.apply(mode="attach", env_ids=...). Delegates to SimulationManager.create_rigid_constraint().

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (Tensor | None) – Target environment indices. None -> all envs.

  • obj_a_cfg (SceneEntityCfg) – SceneEntityCfg pointing at the first RigidObject.

  • obj_b_cfg (SceneEntityCfg) – SceneEntityCfg pointing at the second RigidObject.

  • name (str) – Base constraint name; per-arena names derived by the sim layer.

  • local_frame_a (ndarray | None) – Local joint frame on object A. None -> identity (object A’s origin). Accepts (4,4) or (N,4,4).

  • local_frame_b (ndarray | None) – Local joint frame on object B. None -> computed per env as inv(pose_B) @ pose_A so the constraint welds the objects at their current relative pose. Accepts (4,4) or (N,4,4).

Raises:

RuntimeError – If either entity is not a RigidObject.

Return type:

None

embodichain.lab.gym.envs.managers.events.drop_rigid_object_group_sequentially(env, env_ids, entity_cfg, drop_position=[0.0, 0.0, 1.0], position_range=([-0.1, -0.1, 0.0], [0.1, 0.1, 0.0]), physics_step=2)[source]#

Drop rigid object group from a specified height sequentially in the environment.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | None) – The environment IDs to apply the event.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity to randomize.

  • drop_position (List[float]) – The base position from which to drop the objects. Default is [0.0, 0.0, 1.0].

  • position_range (Tuple[List[float], List[float]]) – The range for randomizing the drop position around the base position.

  • physics_step (int) – The number of physics steps to simulate after dropping the objects. Default is 2.

Return type:

None

class embodichain.lab.gym.envs.managers.events.prepare_extra_attr[source]#

Methods:

__init__(cfg, env)

Initializes the event manager with the given configuration and environment.

__init__(cfg, env)[source]#

Initializes the event manager with the given configuration and environment.

Parameters:
  • cfg (FunctorCfg) – The configuration object for the functor.

  • env (EmbodiedEnv) – The embodied environment instance.

extra_attrs#

A dictionary to hold additional attributes.

Type:

dict

embodichain.lab.gym.envs.managers.events.register_entity_attrs(env, env_ids, entity_cfg, registration='affordance_datas', attrs=[], prefix=True)[source]#

Register the atrributes of an entity to the env.registration dict.

TODO: Currently this method only support 1 env or multi-envs that reset() together,

as it’s behavior is to update a overall dict every time it’s called.

In the future, asynchronously reset mode shall be supported.

Parameters:
  • env (EmbodiedEnv) – The environment the entity is in.

  • env_ids (torch.Tensor | None) – The ids of the envs that the entity should be registered.

  • entity_cfg (SceneEntityCfg) – The config of the entity.

  • attrs (List[str]) – The list of entity attributes that asked to be registered.

  • registration (str, optional) – The env’s registration string where the attributes should be injected to.

embodichain.lab.gym.envs.managers.events.register_info_to_env(env, env_ids, registry, registration='affordance_datas', sim_update=False)[source]#

Register entity information without advancing physics by default.

Parameters:
  • env (EmbodiedEnv) – Environment receiving the registered information.

  • env_ids (Tensor | None) – Environment indices to read. None selects all environments.

  • registry (List[Dict]) – Entity information declarations.

  • registration (str) – Name of the environment’s destination mapping.

  • sim_update (bool) – Explicitly run the legacy 100-step settle before reading. This requires all environments because physics advances the World.

Return type:

None

embodichain.lab.gym.envs.managers.events.remove_rigid_constraint(env, env_ids, name)[source]#

Remove the named constraint for the given env_ids.

Delegates to SimulationManager.remove_rigid_constraint(). Idempotent: warns (via the sim layer) if the constraint is not found.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (Tensor | None) – Target environment indices. None -> all envs.

  • name (str) – Base constraint name to remove.

Return type:

None

class embodichain.lab.gym.envs.managers.events.replace_assets_from_group[source]#

Replace a directory-sampled mesh or a prepared semantic asset combination.

Existing entity_cfg/folder_path configurations keep their directory sampling behavior. Alternatively, supply deployment from asset-batch preparation. That mode jointly replaces the selected rigid bodies, retains the environment and installs the adapted Task Program binding. It requires a cleared B=1 episode boundary. Provider I/O and semantic matching belong to preparation and do not run inside this event.

Methods:

__init__(cfg, env)

Initialize the term.

__init__(cfg, env)[source]#

Initialize the term.

Parameters:
Raises:

ValueError – If the asset is not a RigidObject or an Articulation.

embodichain.lab.gym.envs.managers.events.set_detached_uids_for_env_reset(env, env_ids, uids)[source]#

Set the UIDs of objects that are detached from automatic reset in the environment.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | None) – The environment IDs to apply the event.

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

Return type:

None

embodichain.lab.gym.envs.managers.events.wait_for_dynamic_objects_to_settle(env, env_ids, entity_cfgs=None, linear_velocity_threshold=0.03, angular_velocity_threshold=0.2, min_steps=10, max_steps=240, check_interval_steps=2, required_stable_checks=3, timeout_behavior='warn', allow_partial_envs=False)[source]#

Advance physics until selected dynamic objects remain stationary.

The functor waits at least min_steps and then polls every check_interval_steps. Every selected body in every selected environment must remain below both velocity thresholds for required_stable_checks consecutive polls. It never clears dynamics and never advances beyond max_steps.

When entity_cfgs is None, all dynamic rigid objects, rigid object groups, and non-robot articulations are selected automatically. Static and kinematic entities are ignored during automatic discovery.

Attention

SimulationManager.update() advances the entire vectorized physics world. Partial env_ids are therefore rejected by default. Set allow_partial_envs=True only when advancing non-target environments is acceptable.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (Tensor | Sequence[int] | slice | None) – Target environment IDs. None or slice(None) selects all environments.

  • entity_cfgs (Sequence[SceneEntityCfg] | None) – Explicit settle targets. None discovers all supported dynamic non-robot entities.

  • linear_velocity_threshold (float) – Maximum stable linear speed in meters per second.

  • angular_velocity_threshold (float) – Maximum stable angular speed in radians per second.

  • min_steps (int) – Physics steps to run before the first stability check.

  • max_steps (int) – Maximum total number of physics steps to run.

  • check_interval_steps (int) – Physics steps between stability checks.

  • required_stable_checks (int) – Consecutive stable checks required before return.

  • timeout_behavior (Literal['warn', 'raise']) – "warn" to log and continue or "raise" to raise TimeoutError when max_steps is reached.

  • allow_partial_envs (bool) – Whether to permit a partial environment selection despite whole-world physics advancement.

Raises:
  • IndexError – If an environment ID is outside the valid range.

  • RuntimeError – If a dynamic target has no readable body data.

  • TimeoutError – If objects do not settle and timeout_behavior is "raise".

  • TypeError – If a parameter or entity configuration has the wrong type.

  • ValueError – If parameters, targets, or environment selection are invalid.

Return type:

None

Recording Functions#

Classes:

record_camera_data

Record independent environment videos, committed explicitly before reset.

record_camera_data_async

Persist independent row snapshots through one background video writer.

validation_cameras

This functor creates validation cameras during initialization and captures their data when called.

class embodichain.lab.gym.envs.managers.record.record_camera_data[source]#

Record independent environment videos, committed explicitly before reset.

Frames and provenance are owned per row. Normal commits obey the environment’s EpisodeContext; purpose="diagnostic" labels isolated audit exports. Uncommitted frames are discarded at close, never implicitly published.

Methods:

__init__(cfg, env)

Initialize the functor.

close()

Finalize committed jobs and discard uncommitted frames.

discard_and_clear([env_ids])

Discard selected live frames while preserving committed jobs.

finalize()

Flush committed jobs and discard live frames; never commit on close.

save_and_clear([env_ids, purpose])

Commit selected eligible rows, retaining failed writes for retry.

__init__(cfg, env)[source]#

Initialize the functor.

Parameters:
Raises:

ValueError – If the asset is not a RigidObject or an Articulation.

close()[source]#

Finalize committed jobs and discard uncommitted frames.

Return type:

None

discard_and_clear(env_ids=None)[source]#

Discard selected live frames while preserving committed jobs.

Parameters:

env_ids (Tensor | list[int] | None) – Environment rows to discard; None selects all rows.

Return type:

None

finalize()[source]#

Flush committed jobs and discard live frames; never commit on close.

Return type:

None

save_and_clear(env_ids=None, *, purpose='recording')[source]#

Commit selected eligible rows, retaining failed writes for retry.

Parameters:
  • env_ids (Tensor | list[int] | None) – Environment rows to commit; None selects all rows.

  • purpose (str) – recording obeys eligibility; diagnostic writes explicitly labeled artifacts under a separate diagnostics directory.

Return type:

None

class embodichain.lab.gym.envs.managers.record.record_camera_data_async[source]#

Persist independent row snapshots through one background video writer.

Methods:

__init__(cfg, env)

Initialize the functor.

finalize()

Wait for every committed row; independent rows never wait for a mosaic.

__init__(cfg, env)[source]#

Initialize the functor.

Parameters:
Raises:

ValueError – If the asset is not a RigidObject or an Articulation.

finalize()[source]#

Wait for every committed row; independent rows never wait for a mosaic.

Return type:

None

class embodichain.lab.gym.envs.managers.record.validation_cameras[source]#

This functor creates validation cameras during initialization and captures their data when called. The cameras are created once and reused for subsequent calls.

Methods:

__init__(cfg, env)

Initialize the functor.

__init__(cfg, env)[source]#

Initialize the functor.

Parameters:

Dataset Recording#

Dataset manager for orchestrating dataset collection functors.

Classes:

DatasetManager

Manager for orchestrating dataset collection and saving using functors.

class embodichain.lab.gym.envs.managers.dataset_manager.DatasetManager[source]#

Manager for orchestrating dataset collection and saving using functors.

The dataset manager supports multiple dataset formats through a functor system: - LeRobot format (via LeRobotRecorder) - HDF5 format (via HDF5Recorder) - Zarr format (via ZarrRecorder) - Custom formats (via user-defined functors)

Each functor’s step() method is called once per environment step and handles: - Recording observation-action pairs - Detecting episode completion (dones=True) - Auto-saving completed episodes

Example configuration:
>>> from embodichain.lab.gym.envs.managers.cfg import DatasetFunctorCfg
>>> from embodichain.lab.gym.envs.managers.datasets import LeRobotRecorder
>>>
>>> @configclass
>>> class MyEnvCfg:
>>>     dataset: dict = {
>>>         "lerobot": DatasetFunctorCfg(
>>>             func=LeRobotRecorder,
>>>             save_failed_episodes=True,
>>>             params={
>>>                 "robot_meta": {...},
>>>                 "instruction": {"lang": "pick and place"},
>>>                 "extra": {"scene_type": "kitchen"},
>>>                 "save_path": "/data/datasets"
>>>             }
>>>         )
>>>     }

Methods:

__init__(cfg, env)

Initialize the dataset manager.

apply(mode[, env_ids])

Apply dataset functors for the specified mode.

clear_cache()

Clear cached data from all dataset functors (for online training).

close()

Finalize all dataset functors; repeated calls are safe.

finalize()

Finalize every dataset functor exactly once.

get_cached_data()

Get cached data from all dataset functors (for online training).

get_functor_cfg(functor_name)

Gets the configuration for the specified functor.

reset([env_ids])

Reset all dataset functors.

validate_policy_action(action)

Validate a live policy action for every descriptor-driven recorder.

Attributes:

active_functors

Name of active dataset functors by mode.

available_modes

List of available modes for the dataset manager.

requires_raw_actions

Whether a recorder needs pre-controller policy actions.

save_failed_episodes

Whether any configured dataset recorder should keep failed episodes.

__init__(cfg, env)[source]#

Initialize the dataset manager.

Parameters:
  • cfg (object) – Configuration object containing dataset functor configurations.

  • env (EmbodiedEnv) – The environment instance.

property active_functors: dict[str, list[str]]#

Name of active dataset functors by mode.

The keys are the modes and the values are the names of the dataset functors.

apply(mode, env_ids=None)[source]#

Apply dataset functors for the specified mode.

This method saves completed episodes by reading data from the environment’s episode buffers. It should be called before clearing the buffers during reset.

Parameters:
  • mode (str) – The mode to apply (currently only “save” is supported).

  • env_ids (Union[Sequence[int], Tensor, None]) – The indices of the environments to apply the functor to. Defaults to None, in which case the functor is applied to all environments.

Return type:

None

property available_modes: list[str]#

List of available modes for the dataset manager.

clear_cache()[source]#

Clear cached data from all dataset functors (for online training).

Iterates through all functors and clears their cache if they support online training mode (have clear_cache method).

Return type:

int

Returns:

Total number of cached items cleared across all functors.

close()[source]#

Finalize all dataset functors; repeated calls are safe.

Return type:

Optional[str]

finalize()[source]#

Finalize every dataset functor exactly once.

Finalization is a storage barrier only; individual recorders must not implicitly commit live episode buffers here. All functors are attempted even when one fails, and their failures are reported together.

Return type:

Optional[str]

Returns:

Path to the first finalized dataset, or None if none was returned.

Raises:

RuntimeError – If one or more functors fail to finalize.

get_cached_data()[source]#

Get cached data from all dataset functors (for online training).

Iterates through all functors and collects cached data from those that support online training mode (have get_cached_data method).

Return type:

list[Dict[str, Any]]

Returns:

List of cached data dictionaries from all functors.

get_functor_cfg(functor_name)[source]#

Gets the configuration for the specified functor.

Parameters:

functor_name (str) – The name of the dataset functor.

Return type:

DatasetFunctorCfg

Returns:

The configuration of the dataset functor.

Raises:

ValueError – If the functor name is not found.

property requires_raw_actions: bool#

Whether a recorder needs pre-controller policy actions.

reset(env_ids=None)[source]#

Reset all dataset functors.

Parameters:

env_ids (Union[Sequence[int], Tensor, None]) – The environment ids. Defaults to None.

Return type:

dict[str, float]

Returns:

Empty dict (no logging info).

property save_failed_episodes: bool#

Whether any configured dataset recorder should keep failed episodes.

validate_policy_action(action)[source]#

Validate a live policy action for every descriptor-driven recorder.

Parameters:

action (Tensor) – Flat policy tensor before ActionManager processing.

Return type:

None

Dataset functors for collecting and saving episode data.

Classes:

LeRobotRecorder

Functor for recording episodes in LeRobot format.

class embodichain.lab.gym.envs.managers.datasets.LeRobotRecorder[source]#

Functor for recording episodes in LeRobot format.

This functor handles:

  • Recording observation-action pairs during episodes

  • Converting data to LeRobot format

  • Saving episodes when they complete

Without an action_contract parameter, the recorder preserves the existing joint-action schema. An optional versioned action contract can select a joint or EEF representation and its auxiliary fields explicitly. All representations retain the overall instruction as LeRobot’s task and the current segment instruction as subtask, including fragments.

Methods:

__init__(cfg, env)

Initialize the LeRobot dataset recorder.

close()

Finalize the recorder; repeated calls are safe.

finalize()

Finalize resources without implicitly committing a partial episode.

validate_policy_action(action)

Reject invalid live policy values before term processing or execution.

Attributes:

dataset_path

Path to the dataset directory.

requires_raw_actions

Whether this recorder needs raw policy requests.

__init__(cfg, env)[source]#

Initialize the LeRobot dataset recorder.

Parameters:
  • cfg (DatasetFunctorCfg) –

    Functor configuration containing params: - save_path: Root directory for saving datasets - robot_meta: Robot metadata for dataset - instruction: Optional legacy fallback; task/episode language takes precedence - extra: Optional extra metadata - use_videos: Whether to save videos - image_writer_threads: Number of threads for image writing - image_writer_processes: Number of processes for image writing - image_compress_level: Optional PNG compression level (0-9)

    for image features; video temporary PNGs keep LeRobot’s default level 1

    • metadata_buffer_size: Positive LeRobot episode metadata buffer size; default 1. Larger buffers increase data at risk until the SDK is finalized.

  • env (EmbodiedEnv) – The environment instance

close()[source]#

Finalize the recorder; repeated calls are safe.

Return type:

Optional[str]

property dataset_path: str#

Path to the dataset directory.

finalize()[source]#

Finalize resources without implicitly committing a partial episode.

Episodes are committed only when __call__() is invoked by an explicit reset(save_data=True). Closing the environment therefore leaves any still-live rollout buffer uncommitted.

Return type:

Optional[str]

Returns:

The finalized dataset path, or None when no dataset exists.

Raises:

RuntimeError – If one or more dataset resources cannot be finalized.

property requires_raw_actions: bool#

Whether this recorder needs raw policy requests.

validate_policy_action(action)[source]#

Reject invalid live policy values before term processing or execution.

Parameters:

action (Tensor) – Flat policy tensor for all vector environments.

Raises:
  • TypeError – If the action is not a floating tensor.

  • ValueError – If its layout or descriptor-owned values are invalid.

Return type:

None

Asynchronous LeRobot recorder for parallel environments.

This module provides AsyncLeRobotRecorder, which decouples episode saving from the simulation loop. It is the recommended recorder when running many parallel environments that complete episodes together: instead of blocking env.reset() while episodes are converted and flushed to disk, the completed episode buffers are cloned and handed to a background worker thread, so the simulator can keep stepping.

Classes:

AsyncLeRobotRecorder

LeRobot recorder that saves episodes on a background thread.

class embodichain.lab.gym.envs.managers.async_datasets.AsyncLeRobotRecorder[source]#

LeRobot recorder that saves episodes on a background thread.

Drop-in replacement for LeRobotRecorder selected via the dataset config "func": "AsyncLeRobotRecorder". It shares the same on-disk format and feature-building logic; only the timing of the save differs.

Why this helps parallel environments#

In the synchronous recorder, LeRobotRecorder.__call__() runs inside env.reset() (via DatasetManager.apply) and blocks the simulator while it iterates add_frame + save_episode for every finished env. With num_envs=N all finishing at once, the sim stalls for the sum of all episodes’ save time every reset.

This recorder instead, on each apply:

  1. Reads each env’s rollout-buffer slice (obs/actions).

  2. Clones the slice to CPU (detached from the live buffer).

  3. Clones frame annotations and episode/segment metadata with the payload.

  4. Pushes the detached payload onto a queue.

  5. Returns immediately - the sim is free to reset and keep stepping.

A single daemon worker thread drains the queue and runs the standard LeRobotRecorder._persist_episode_payload() on each cloned payload.

Correctness#

  • No concurrent dataset access. LeRobotDataset is not thread-safe; only the worker thread ever calls add_frame / save_episode / mutates curr_episode. The main thread only enqueues and, at close, drains.

  • No buffer race. The slice is cloned in the caller thread before the buffer is cleared on reset, so the worker never reads memory that the sim is overwriting.

  • Ordering. A single worker preserves FIFO episode order, so episode_index assignment is deterministic.

  • Drain on close. finalize() joins the worker before the parent flushes the image writer and finalizes the dataset.

Note

The clone copies each episode’s camera frames into host RAM. The default queue is unbounded. A positive async_queue_maxsize limits waiting payloads; the active worker payload, caller’s current clone, rollout buffer and image-writer queue consume additional memory. When the queue is full, committing an episode blocks until space is available, allowing storage to catch up without dropping data.

type cfg:

DatasetFunctorCfg

param cfg:

DatasetFunctorCfg with the same params as LeRobotRecorder. The image_writer_threads / image_writer_processes params are honored and combine with the background worker (two levels of async: episode conversion off the sim thread, PNG writes off the worker thread). async_queue_maxsize optionally bounds the number of committed payloads waiting for the worker; zero keeps the historical unbounded queue behavior.

type env:

EmbodiedEnv

param env:

The environment instance.

Methods:

__init__(cfg, env)

Initialize the LeRobot dataset recorder.

finalize()

Drain committed writes, finalize storage, and surface all failures.

Attributes:

async_stats

Return a snapshot of queue depth and producer waiting time.

__init__(cfg, env)[source]#

Initialize the LeRobot dataset recorder.

Parameters:
  • cfg (DatasetFunctorCfg) –

    Functor configuration containing params: - save_path: Root directory for saving datasets - robot_meta: Robot metadata for dataset - instruction: Optional legacy fallback; task/episode language takes precedence - extra: Optional extra metadata - use_videos: Whether to save videos - image_writer_threads: Number of threads for image writing - image_writer_processes: Number of processes for image writing - image_compress_level: Optional PNG compression level (0-9)

    for image features; video temporary PNGs keep LeRobot’s default level 1

    • metadata_buffer_size: Positive LeRobot episode metadata buffer size; default 1. Larger buffers increase data at risk until the SDK is finalized.

  • env (EmbodiedEnv) – The environment instance

property async_stats: dict[str, float | int]#

Return a snapshot of queue depth and producer waiting time.

Queue sizes and their sampled peak are approximate. Backpressure counts attempts that found the queue full and measures only time waiting to enqueue, excluding payload cloning and finalization.

Returns:

Queue size/limit/peak, blocked enqueue count and waiting seconds.

finalize()[source]#

Drain committed writes, finalize storage, and surface all failures.

The queue is a durability barrier for episodes explicitly committed by reset(save_data=True). A live rollout that was never enqueued is not saved during close.

Return type:

Optional[str]

Returns:

The finalized dataset path.

Raises:

RuntimeError – If any queued episode or dataset resource failed.

Randomization#

Domain-randomization event functors (physics, visual, spatial, geometry).

Implemented as event functors registered via EventCfg and dispatched by EventManager at startup / reset / interval modes.

Submodules

Physics#

Classes:

push_articulation_by_setting_velocity

Apply independently scheduled root-velocity disturbances.

Functions:

randomize_articulation_mass(env, env_ids, ...)

Randomize the mass of articulation links in the environment.

randomize_rigid_object_center_of_mass(env, ...)

Randomize the center of mass of rigid objects in the environment.

randomize_rigid_object_mass(env, env_ids, ...)

Randomize the mass of rigid objects in the environment.

class embodichain.lab.gym.envs.managers.randomization.physics.push_articulation_by_setting_velocity[source]#

Apply independently scheduled root-velocity disturbances.

Methods:

__init__(cfg, env)

Initialize the manager term from the environment and configuration.

reset([env_ids])

Resample disturbance timers for selected environments.

__init__(cfg, env)[source]#

Initialize the manager term from the environment and configuration.

Parameters:
  • cfg (FunctorCfg) – Manager configuration containing the term parameters.

  • env (EmbodiedEnv) – Environment providing the task state and robot.

reset(env_ids=None)[source]#

Resample disturbance timers for selected environments.

Parameters:

env_ids (Optional[Sequence[int]]) – Selected environment rows; None selects all rows when supported.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.physics.randomize_articulation_mass(env, env_ids, entity_cfg, mass_range, link_names=None, relative=False, recompute_inertia=True, min_mass=1e-06)[source]#

Randomize the mass of articulation links in the environment.

Uses regular expression matching to select which links to randomize.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | list[int]) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration for the scene entity.

  • mass_range (tuple[float, float] | dict[str, tuple[float, float]]) – The range (min, max) to sample the mass from. If a dict, keys are link names and values are per-link mass ranges. When a dict is provided, link_names is ignored and the dict keys are used instead.

  • link_names (str | list[str] | None) – A regex pattern or list of regex patterns to match link names. If None, all links are randomized. Ignored when mass_range is a dict. Defaults to None.

  • relative (bool) – Whether to apply the mass change relative to the initial mass. Defaults to False.

  • recompute_inertia (bool) – Whether to scale initialization-time inertia by the sampled mass ratio. Defaults to True.

  • min_mass (float) – Minimum allowed sampled mass. Defaults to 1e-6.

Raises:

ValueError – If min_mass or an initialization-time link mass is not positive.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.physics.randomize_rigid_object_center_of_mass(env, env_ids, entity_cfg, com_pos_offset_range)[source]#

Randomize the center of mass of rigid objects in the environment.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | list[int]) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration for the scene entity.

  • com_pos_offset_range (tuple[list[float], list[float]]) – The range (min, max) to sample the center of mass offset from.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.physics.randomize_rigid_object_mass(env, env_ids, entity_cfg, mass_range, relative=False, recompute_inertia=True, min_mass=1e-06)[source]#

Randomize the mass of rigid objects in the environment.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | list[int]) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration for the scene entity.

  • mass_range (tuple[float, float]) – The range (min, max) to sample the mass from.

  • relative (bool) – Whether to apply the mass change relative to the initial mass. Defaults to False.

  • recompute_inertia (bool) – Whether to scale the initial inertia by the sampled mass ratio. Defaults to True.

  • min_mass (float) – Minimum allowed sampled mass. Defaults to 1e-6.

Raises:

ValueError – If min_mass is not positive or an initial mass is not positive.

Return type:

None

Visual#

Functions:

randomize_camera_extrinsics(env, env_ids, ...)

Randomize camera extrinsic properties (position and orientation).

randomize_camera_intrinsics(env, env_ids, ...)

Randomize camera intrinsic properties by adding, scaling, or setting random values.

randomize_emission_light(env, env_ids[, ...])

Randomize emission light properties by adding, scaling, or setting random values.

randomize_light(env, env_ids, entity_cfg[, ...])

Randomize light properties by adding, scaling, or setting random values.

set_rigid_object_group_visual_material(env, ...)

Set a rigid object group's visual material (deterministic, non-random).

set_rigid_object_visual_material(env, ...)

Set a rigid object's visual material (deterministic, non-random).

Classes:

randomize_indirect_lighting

Randomize the environment's indirect (IBL) lighting or emissive light.

randomize_visual_material

Randomize the visual material properties of a RigidObject or Articulation.

embodichain.lab.gym.envs.managers.randomization.visual.randomize_camera_extrinsics(env, env_ids, entity_cfg, pos_range=None, euler_range=None, eye_range=None, target_range=None, up_range=None)[source]#

Randomize camera extrinsic properties (position and orientation).

Behavior: - Pose extrinsics use pos_range/euler_range whether the camera is

mounted to a parent or placed directly in the arena.

  • Look-at extrinsics use eye_range/target_range/up_range and recompute the camera pose with look_at.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (Union[Sequence[int], Tensor, None]) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity to randomize.

  • pos_range (tuple[list[float], list[float]] | None) – Position perturbation range (pose mode).

  • euler_range (tuple[list[float], list[float]] | None) – Euler angle perturbation range (pose mode).

  • eye_range (tuple[list[float], list[float]] | None) – Eye position perturbation range (look_at mode).

  • target_range (tuple[list[float], list[float]] | None) – Target position perturbation range (look_at mode).

  • up_range (tuple[list[float], list[float]] | None) – Up vector perturbation range (look_at mode).

Return type:

None

Ranges that do not apply to the camera’s extrinsics mode are ignored; a warning is logged once per camera so that such config mistakes are visible.

embodichain.lab.gym.envs.managers.randomization.visual.randomize_camera_intrinsics(env, env_ids, entity_cfg, focal_x_range=None, focal_y_range=None, cx_range=None, cy_range=None)[source]#

Randomize camera intrinsic properties by adding, scaling, or setting random values.

This function allows randomizing camera intrinsic parameters in the scene. The function samples random values from the given distribution parameters and adds, scales, or sets the values into the physics simulation based on the operation.

The distribution parameters are tuples of two elements each, representing the lower and upper bounds of the distribution for the focal length (fx, fy) and principal point (cx, cy) components of the camera intrinsics. The function samples random values for each component independently.

Attention

This function applies the same intrinsic properties for all the environments.

focal_x_range and focal_y_range are values added to the camera’s current fx and fy values. focal_xy_range is a combined range for both fx and fy, where the range is specified as [[fx_min, fy_min], [fx_max, fy_max]]. cx_range and cy_range are values added to the camera’s current cx and cy values.

Tip

This function uses CPU tensors to assign camera intrinsic properties.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (Union[torch.Tensor, None]) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity to randomize.

  • focal_x_range (tuple[float, float] | None) – The range for the focal length x randomization.

  • focal_y_range (tuple[float, float] | None) – The range for the focal length y randomization.

  • cx_range (tuple[float, float] | None) – The range for the principal point x randomization.

  • cy_range (tuple[float, float] | None) – The range for the principal point y randomization.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.visual.randomize_emission_light(env, env_ids, color_range=None, intensity_range=None)[source]#

Randomize emission light properties by adding, scaling, or setting random values.

This function allows randomizing emission light properties in the scene. The function samples random values from the given distribution parameters and adds, scales, or sets the values into the physics simulation based on the operation.

The distribution parameters are lists of two elements each, representing the lower and upper bounds of the distribution for the r, g, b components of the light color and intensity. The function samples random values for each component independently. :rtype: None

Attention

This function applied the same emission light properties for all the environments.

color_range is the absolute r, g, b value set on the emission light. intensity_range is the absolute intensity value set on the emission light.

class embodichain.lab.gym.envs.managers.randomization.visual.randomize_indirect_lighting[source]#

Randomize the environment’s indirect (IBL) lighting or emissive light.

This functor operates in one of two mutually exclusive modes:

  • HDR mode — path is provided. A random .hdr file is chosen from the folder on every call and applied via set_indirect_lighting().

  • Emissive mode — emissive_color_range and/or emissive_intensity_range are provided. The emissive light color and intensity are sampled uniformly on every call and applied via set_emission_light().

Providing both path and emissive parameters simultaneously is an error.

Attention

This functor applies the same lighting to all environments.

Tip

The path parameter is resolved via get_data_path(), so it supports absolute paths, data-root-relative paths, and dataset-class paths (e.g. "EnvMapHDR").

emissive_color_range is a pair of [r, g, b] lists representing the lower and upper bounds for sampling the emissive color, e.g. [[0.8, 0.8, 0.8], [1.0, 1.0, 1.0]].

emissive_intensity_range is a [min, max] pair for the emissive intensity scalar, e.g. [80.0, 150.0].

Methods:

__init__(cfg, env)

Initialize the functor.

__init__(cfg, env)[source]#

Initialize the functor.

Parameters:
  • cfg (FunctorCfg) –

    The configuration of the functor.

    • HDR mode: set params["path"] to a folder of .hdr files.

    • Emissive mode: set params["emissive_color_range"] (pair of RGB lists) and/or params["emissive_intensity_range"] (pair of floats).

  • env (EmbodiedEnv) – The environment instance.

Raises:

ValueError – If both HDR and emissive params are provided, or if neither is provided.

embodichain.lab.gym.envs.managers.randomization.visual.randomize_light(env, env_ids, entity_cfg, position_range=None, color_range=None, intensity_range=None, direction_range=None)[source]#

Randomize light properties by adding, scaling, or setting random values.

This function allows randomizing light properties in the scene. The function samples random values from the given distribution parameters and adds, scales, or sets the values into the physics simulation based on the operation.

The distribution parameters are lists of two elements each, representing the lower and upper bounds of the distribution for the x, y, and z components of the light properties. The function samples random values for each component independently.

Attention

This function applied the same light properties for all the environments.

position_range is the x, y, z value added into light’s cfg.init_pos. color_range is the absolute r, g, b value set to the light object. intensity_range is the value added into light’s cfg.intensity. direction_range is the x, y, z value added into light’s cfg.direction. (Only applicable for "sun", "direction", "spot", "rect", and "mesh" light types.)

Tip

This function uses CPU tensors to assign light properties.

Warning

position_range is ignored for global scene lights ("sun", "direction") because they are infinite-distance lights with no meaningful position. Use direction_range instead for these light types.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (Union[torch.Tensor, None]) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity to randomize.

  • position_range (tuple[list[float], list[float]] | None) – The range for the position randomization.

  • color_range (tuple[list[float], list[float]] | None) – The range for the color randomization.

  • intensity_range (tuple[float, float] | None) – The range for the intensity randomization.

  • direction_range (tuple[list[float], list[float]] | None) – The range for the direction randomization. Only applicable for directional light types ("sun", "direction", "spot", "rect", "mesh").

Return type:

None

class embodichain.lab.gym.envs.managers.randomization.visual.randomize_visual_material[source]#

Randomize the visual material properties of a RigidObject or Articulation.

Supported properties are base color, base-color texture, metallic factor, roughness, and index of refraction. Textures from texture_path are preloaded during initialization. The default ground plane can be selected with entity_cfg.uid="default_plane".

By default, the functor retains the asset’s original material instances and creates working instances from their existing templates. Set fallback_to_new=True to force the legacy new-material path.

Methods:

__init__(cfg, env)

Initialize the term.

gen_random_base_color_texture(width, height)

Generate a random base color texture.

__init__(cfg, env)[source]#

Initialize the term.

Parameters:
Raises:

ValueError – If the asset is not a RigidObject or an Articulation.

static gen_random_base_color_texture(width, height)[source]#

Generate a random base color texture.

Parameters:
  • width (int) – The width of the texture.

  • height (int) – The height of the texture.

Return type:

Tensor

Returns:

A torch tensor representing the random base color texture with shape (height, width, 4).

embodichain.lab.gym.envs.managers.randomization.visual.set_rigid_object_group_visual_material(env, env_ids, entity_cfg, mat_cfg)[source]#

Set a rigid object group’s visual material (deterministic, non-random).

This helper exists to support configs that want fixed colors/materials during reset.

Parameters:
  • env (EmbodiedEnv) – Environment instance.

  • env_ids (Tensor | None) – Target env ids. If None, applies to all envs.

  • entity_cfg (SceneEntityCfg) – Scene entity config (must point to a rigid object).

  • mat_cfg (Union[VisualMaterialCfg, Dict]) – Visual material configuration. Can be a VisualMaterialCfg object or a dict. If a dict is provided, it will be converted to VisualMaterialCfg using from_dict(). If uid is not specified in mat_cfg, it will default to “{entity_uid}_mat”.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.visual.set_rigid_object_visual_material(env, env_ids, entity_cfg, mat_cfg)[source]#

Set a rigid object’s visual material (deterministic, non-random).

This helper exists to support configs that want fixed colors/materials during reset.

Parameters:
  • env (EmbodiedEnv) – Environment instance.

  • env_ids (Optional[Tensor]) – Target env ids. If None, applies to all envs.

  • entity_cfg (SceneEntityCfg) – Scene entity config (must point to a rigid object).

  • mat_cfg (Union[VisualMaterialCfg, Dict]) – Visual material configuration. Can be a VisualMaterialCfg object or a dict. If a dict is provided, it will be converted to VisualMaterialCfg using from_dict(). If uid is not specified in mat_cfg, it will default to “{entity_uid}_mat”.

Return type:

None

Spatial#

Classes:

ExpandSceneLayout

Apply bounded online or saved layouts through an ordinary reset event.

planner_grid_cell_sampler

Sample grid cells for object placement without replacement.

randomize_anchor_height

Randomize the height of an anchor object and shift other objects by the same delta.

sample_rigid_object_pose_from_workspace

Place a rigid object at a position sampled from a robot workspace.

Functions:

get_random_pose(init_pos, init_rot[, ...])

Generate a random pose based on the initial position and rotation.

randomize_articulation_root_pose(env, ...[, ...])

Randomize the root pose of an articulation in the environment.

randomize_rigid_object_pose(env, env_ids, ...)

Randomize the pose of a rigid object in the environment.

randomize_robot_eef_pose(env, env_ids, ...)

Randomize the initial end-effector pose of a robot in the environment.

randomize_robot_qpos(env, env_ids, entity_cfg)

Randomize the initial joint positions of a robot in the environment.

randomize_target_pose(env, env_ids, ...[, ...])

Randomize a virtual target pose and store in env state.

class embodichain.lab.gym.envs.managers.randomization.spatial.ExpandSceneLayout[source]#

Apply bounded online or saved layouts through an ordinary reset event.

Parameters:
  • cfg (FunctorCfg) – EventCfg with support entity_cfg, complete entity_cfgs, parent_by_uid, source and policy in params. Source kind is online or saved; saved sources require layout_dir. Optional candidate_count/max_attempts control each bounded refill.

  • env (EmbodiedEnv) – Caller-owned manager environment; no Task Program is required.

Geometry is prepared once. Ordinary reset retains prepared geometry and unused candidates. EventManager explicit reseeding rewinds candidate and replay cursors. This functor does not step physics or certify task success. Provenance and recording eligibility belong to env.episode_context. Each event reports under events/<configuration name>.

Methods:

__init__(cfg, env)

Initialize the functor.

reset([env_ids])

Retain prepared geometry and streams; the host owns episode results.

__init__(cfg, env)[source]#

Initialize the functor.

Parameters:
reset(env_ids=None)[source]#

Retain prepared geometry and streams; the host owns episode results.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.spatial.get_random_pose(init_pos, init_rot, position_range=None, rotation_range=None, relative_position=True, relative_rotation=False)[source]#

Generate a random pose based on the initial position and rotation.

Parameters:
  • init_pos (torch.Tensor) – The initial position tensor of shape (num_instance, 3).

  • init_rot (torch.Tensor) – The initial rotation tensor of shape (num_instance, 3, 3).

  • position_range (tuple[list[float], list[float]] | None) – The range for the position randomization.

  • rotation_range (tuple[list[float], list[float]] | None) – The range for the rotation randomization. The rotation is represented as Euler angles (roll, pitch, yaw) in degree.

  • relative_position (bool) – Whether to randomize the position relative to the initial position. Default is True.

  • relative_rotation (bool) – Whether to randomize the rotation relative to the initial rotation. Default is False.

Returns:

The generated random pose tensor of shape (num_instance, 4, 4).

Return type:

torch.Tensor

class embodichain.lab.gym.envs.managers.randomization.spatial.planner_grid_cell_sampler[source]#

Sample grid cells for object placement without replacement.

This functor divides a planar region into a regular 2D grid and samples cells to place objects. Each sampled cell will be marked as occupied and will not be resampled until the grid is reset.

The sampler places objects at the center of selected grid cells, with the z-position set to a reference height.

Methods:

__init__(cfg, env)

Initialize the GridCellSampler functor.

reset([env_ids])

Reset the grid sampling state.

__init__(cfg, env)[source]#

Initialize the GridCellSampler functor.

Parameters:
reset(env_ids=None)[source]#

Reset the grid sampling state.

Parameters:

env_ids (Tensor | None) – The environment IDs to reset. If None, resets all environments.

Return type:

None

class embodichain.lab.gym.envs.managers.randomization.spatial.randomize_anchor_height[source]#

Randomize the height of an anchor object and shift other objects by the same delta.

This functor samples a per-environment height delta, moves the anchor object relative to its configured initial position, and adds the same delta to the Z component of every other included object while preserving XY and rotation.

The functor is configured through FunctorCfg parameters, following the same pattern as planner_grid_cell_sampler.

Methods:

__init__(cfg, env)

Initialize the functor.

__init__(cfg, env)[source]#

Initialize the functor.

Parameters:
embodichain.lab.gym.envs.managers.randomization.spatial.randomize_articulation_root_pose(env, env_ids, entity_cfg, position_range=None, rotation_range=None, relative_position=True, relative_rotation=False, physics_update_step=-1)[source]#

Randomize the root pose of an articulation in the environment.

This function randomizes the position and/or rotation of an articulation’s root link. The articulation’s root is the base frame that all other links are attached to.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | None) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity to randomize.

  • position_range (tuple[list[float], list[float]] | None) – The range for the position randomization. Format: [[x_min, y_min, z_min], [x_max, y_max, z_max]].

  • rotation_range (tuple[list[float], list[float]] | None) – The range for the rotation randomization. The rotation is represented as Euler angles (roll, pitch, yaw) in degrees.

  • relative_position (bool) – Whether to randomize the position relative to the articulation’s initial position. Default is True.

  • relative_rotation (bool) – Whether to randomize the rotation relative to the articulation’s initial rotation. Default is False.

  • physics_update_step (int) – The number of physics update steps to apply after randomization. Default is -1 (no update).

Return type:

None

Note

This function is similar to randomize_rigid_object_pose() but operates on articulations (multi-link rigid body systems) rather than single rigid objects.

embodichain.lab.gym.envs.managers.randomization.spatial.randomize_rigid_object_pose(env, env_ids, entity_cfg, position_range=None, rotation_range=None, relative_position=True, relative_rotation=False, physics_update_step=-1)[source]#

Randomize the pose of a rigid object in the environment.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | None) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity to randomize.

  • position_range (tuple[list[float], list[float]] | None) – The range for the position randomization.

  • rotation_range (tuple[list[float], list[float]] | None) – The range for the rotation randomization. The rotation is represented as Euler angles (roll, pitch, yaw) in degree.

  • relative_position (bool) – Whether to randomize the position relative to the object’s initial position. Default is True.

  • relative_rotation (bool) – Whether to randomize the rotation relative to the object’s initial rotation. Default is False.

  • physics_update_step (int) – The number of physics update steps to apply after randomization. Default is -1 (no update).

Return type:

None

embodichain.lab.gym.envs.managers.randomization.spatial.randomize_robot_eef_pose(env, env_ids, entity_cfg, position_range=None, rotation_range=None)[source]#

Randomize the initial end-effector pose of a robot in the environment.

Note

  • The position and rotation are performed randomization in a relative manner.

  • The current state of eef pose is computed based on the current joint positions of the robot.

  • This restores selected state and publishes transforms without stepping physics.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | None) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity to randomize.

  • position_range (tuple[list[float], list[float]] | None) – The range for the position randomization.

  • rotation_range (tuple[list[float], list[float]] | None) – The range for the rotation randomization. The rotation is represented as Euler angles (roll, pitch, yaw) in degree.

Raises:

ValueError – No named control parts were supplied for end-effector selection.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.spatial.randomize_robot_qpos(env, env_ids, entity_cfg, qpos_range=None, relative_qpos=True, joint_ids=None)[source]#

Randomize the initial joint positions of a robot in the environment.

Restores current positions and drive targets, clears selected dynamics, and publishes transforms without advancing physics.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | None) – The environment IDs to apply the randomization.

  • entity_cfg (SceneEntityCfg) – The configuration of the scene entity to randomize.

  • qpos_range (tuple[list[float], list[float]] | None) – The range for the joint position randomization.

  • relative_qpos (bool) – Whether to randomize the joint positions relative to the current joint positions. Default is True.

  • joint_ids (List[int] | None) – The list of joint IDs to randomize. If None, all joints will be randomized.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.spatial.randomize_target_pose(env, env_ids, position_range, rotation_range=None, relative_position=False, relative_rotation=False, reference_entity_cfg=None, store_key='target_pose')[source]#

Randomize a virtual target pose and store in env state.

This function generates random target poses without requiring a physical object in the scene. The generated poses are stored as a public attribute in env for use by observations and rewards.

Parameters:
  • env (EmbodiedEnv) – The environment instance.

  • env_ids (torch.Tensor | None) – The environment IDs to apply the randomization.

  • position_range (tuple[list[float], list[float]]) – The range for the position randomization.

  • rotation_range (tuple[list[float], list[float]] | None) – The range for the rotation randomization. The rotation is represented as Euler angles (roll, pitch, yaw) in degree.

  • relative_position (bool) – Whether to randomize the position relative to a reference entity. Default is False.

  • relative_rotation (bool) – Whether to randomize the rotation relative to a reference entity. Default is False.

  • reference_entity_cfg (SceneEntityCfg | None) – The reference entity for relative randomization. If None and relative mode is True, uses world origin.

  • store_key (str) – The key to store the target pose in env state. Default is “target_pose”. The pose will be stored as a public attribute env.{store_key}.

Return type:

None

class embodichain.lab.gym.envs.managers.randomization.spatial.sample_rigid_object_pose_from_workspace[source]#

Place a rigid object at a position sampled from a robot workspace.

The workspace supplies kinematically reachable end-effector samples. This functor consumes their Cartesian positions for object placement while leaving scene collision and motion-planning validation to task-specific logic.

Methods:

__init__(cfg, env)

Initialize the workspace object sampler.

__init__(cfg, env)[source]#

Initialize the workspace object sampler.

Parameters:

Geometry#

Functions:

randomize_rigid_object_body_scale(env, ...)

Deprecated.

randomize_rigid_object_scale(env, env_ids, ...)

Randomize a rigid object's body scale factors (multiplicative, not absolute size).

randomize_rigid_objects_scale(env, env_ids, ...)

Randomize body scale factors for multiple rigid objects.

embodichain.lab.gym.envs.managers.randomization.geometry.randomize_rigid_object_body_scale(env, env_ids, entity_cfg, scale_range=None, same_scale_all_axes=True)[source]#

Deprecated. Use randomize_rigid_object_scale + scale_factor_range.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.geometry.randomize_rigid_object_scale(env, env_ids, entity_cfg, scale_factor_range=None, same_scale_all_axes=True)[source]#

Randomize a rigid object’s body scale factors (multiplicative, not absolute size).

Parameters:
  • env (EmbodiedEnv) – Environment instance.

  • env_ids (Optional[Tensor]) – Target env ids. If None, applies to all envs.

  • entity_cfg (SceneEntityCfg) – Scene entity config of the rigid object.

  • scale_factor_range (tuple[list[float], list[float]] | None) – If same_scale_all_axes is True, should be [[s_min], [s_max]]. Otherwise [[sx_min, sy_min, sz_min], [sx_max, sy_max, sz_max]].

  • same_scale_all_axes (bool) – Whether to use same factor on x/y/z.

Return type:

None

embodichain.lab.gym.envs.managers.randomization.geometry.randomize_rigid_objects_scale(env, env_ids, entity_cfgs, scale_factor_range=None, same_scale_all_axes=True, shared_sample=False)[source]#

Randomize body scale factors for multiple rigid objects.

Parameters:
  • env (EmbodiedEnv) – Environment instance.

  • env_ids (Optional[Tensor]) – Target env ids. If None, applies to all envs.

  • entity_cfgs (List[SceneEntityCfg]) – List of scene entity configs (rigid objects).

  • scale_factor_range (tuple[list[float], list[float]] | None) – Scale factor sampling range.

  • same_scale_all_axes (bool) – Whether to use same factor on x/y/z.

  • shared_sample (bool) – If True, sample one scale per-env and apply to all objects (sync). If False, each object samples its own scales independently.

Return type:

None