# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
from __future__ import annotations
import os
import random
from collections.abc import Mapping, Sequence
from copy import deepcopy
from typing import TYPE_CHECKING, Any, Dict, List, Tuple
import numpy as np
import torch
from embodichain.lab.sim.objects import (
Light,
RigidObject,
RigidObjectGroup,
Articulation,
Robot,
)
from embodichain.lab.sim.cfg import RigidObjectCfg, ArticulationCfg, RigidConstraintCfg
from embodichain.lab.sim.shapes import MeshCfg
from embodichain.lab.gym.envs.managers._event_functors import (
wait_for_dynamic_objects_to_settle,
)
from embodichain.lab.gym.envs.managers.cfg import SceneEntityCfg
from embodichain.lab.gym.envs.managers import Functor, FunctorCfg
from embodichain.utils.module_utils import find_function_from_modules
from embodichain.utils.string import remove_regex_chars, resolve_matching_names
from embodichain.utils.file import get_all_files_in_directory
from embodichain.utils.math import (
sample_uniform,
pose_inv,
xyz_quat_to_4x4_matrix,
trans_matrix_to_xyz_quat,
)
from embodichain.utils import logger
from embodichain.data import get_data_path
if TYPE_CHECKING:
from embodichain.lab.gym.envs import EmbodiedEnv
from embodichain.lab.task_program.integrations.asset_materialization import (
MaterializedAssetDeployment,
)
__all__ = [
"apply_expansion_profile_reset",
"replace_assets_from_group",
"prepare_extra_attr",
"register_entity_attrs",
"register_entity_pose",
"register_info_to_env",
"resolve_uids",
"resolve_dict",
"get_pose",
"drop_rigid_object_group_sequentially",
"wait_for_dynamic_objects_to_settle",
"set_detached_uids_for_env_reset",
"create_rigid_constraint",
"remove_rigid_constraint",
]
[docs]
def apply_expansion_profile_reset(
env: EmbodiedEnv,
env_ids: torch.Tensor | Sequence[int] | slice | None,
*,
cube_pose: torch.Tensor | None = None,
visual_registry: Any | None = None,
visual_assignments: Mapping[int, str] | None = None,
visual_seed: int = 0,
) -> None:
"""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.
Args:
env: Environment whose selected rows should be prepared.
env_ids: Environment rows selected by the reset lifecycle.
cube_pose: Optional ``(num_envs, 7)`` pose tensor in position/xyzw form.
visual_registry: Optional registry exposing ``apply_to_environment``.
visual_assignments: Row-to-profile assignments for the visual registry.
visual_seed: Base seed passed to the visual registry.
"""
if cube_pose is None and visual_registry is None:
return
num_envs = int(env.num_envs)
if env_ids is None:
row_ids = torch.arange(num_envs, dtype=torch.long)
elif isinstance(env_ids, slice):
start, stop, step = env_ids.indices(num_envs)
row_ids = torch.arange(start, stop, step, dtype=torch.long)
else:
row_ids = torch.as_tensor(env_ids, dtype=torch.long).flatten()
if row_ids.numel() == 0:
return
if torch.any(row_ids < 0) or torch.any(row_ids >= num_envs):
raise ValueError("expansion reset env_ids are outside the environment range")
if cube_pose is not None:
if cube_pose.ndim != 2 or cube_pose.shape[1] != 7:
raise ValueError("cube_pose must have shape (N, 7)")
if cube_pose.shape[0] == num_envs:
selected_pose = cube_pose.index_select(
0, row_ids.to(device=cube_pose.device)
)
elif cube_pose.shape[0] == row_ids.numel():
selected_pose = cube_pose
else:
raise ValueError("cube_pose rows must match num_envs or selected env_ids")
cube = env.sim.get_rigid_object("cube")
if cube is None:
raise ValueError("expansion reset requires a rigid object named 'cube'")
env_device = getattr(env, "device", None)
if not isinstance(env_device, (torch.device, str)):
env_device = selected_pose.device
pose_env_ids = row_ids.to(device=env_device)
cube.set_local_pose(selected_pose, env_ids=pose_env_ids)
if visual_registry is not None:
if visual_assignments is None:
raise ValueError("visual_assignments are required with visual_registry")
selected_ids = {int(item) for item in row_ids.tolist()}
assignments = {
int(env_id): profile_id
for env_id, profile_id in visual_assignments.items()
if int(env_id) in selected_ids
}
if not assignments:
raise ValueError("visual_assignments contain no selected environment rows")
visual_registry.apply_to_environment(
env,
assignments,
seed=visual_seed,
)
[docs]
class replace_assets_from_group(Functor):
"""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.
"""
[docs]
def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
"""Initialize the term.
Args:
cfg: The configuration of the functor.
env: The environment instance.
Raises:
ValueError: If the asset is not a RigidObject or an Articulation.
"""
super().__init__(cfg, env)
self._prepared_mode = "deployment" in cfg.params
if self._prepared_mode:
if (
cfg.params.get("entity_cfg") is not None
or cfg.params.get("folder_path") is not None
):
raise ValueError(
"Choose deployment or entity_cfg/folder_path asset replacement"
)
return
# extract the used quantities (to enable type-hinting)
entity_cfg: SceneEntityCfg = cfg.params["entity_cfg"]
asset = env.sim.get_asset(entity_cfg.uid)
if asset is None:
logger.log_error(
f"Asset with UID '{entity_cfg.uid}' not found in the simulation."
)
if (
isinstance(asset, RigidObject)
and isinstance(asset.cfg.shape, MeshCfg) is False
):
logger.log_error(
"Only mesh-based RigidObject assets are supported for replacement."
)
self.asset_cfg = asset.cfg
self.asset_type = type(asset)
if isinstance(asset, Articulation):
logger.log_error("Replacing articulation assets is not supported yet.")
self._asset_group_path: list[str] = []
# The following block of code only handle rigid object assets.
# If we want to support articulation assets, the group path format
# should be changed into list of folder (each folder contains a urdf file
# and its associated resources)
folder_path = cfg.params.get("folder_path", None)
if folder_path is None:
logger.log_error(
"folder_path must be specified in the functor configuration."
)
if folder_path.endswith("/") is False:
folder_path, patterns = os.path.split(folder_path)
# remove regular expression from patterns
patterns = remove_regex_chars(patterns)
self._full_path = get_data_path(f"{folder_path}/")
self._asset_group_path = sorted(
get_all_files_in_directory(self._full_path, patterns=patterns)
)
else:
self._full_path = get_data_path(folder_path)
self._asset_group_path = sorted(get_all_files_in_directory(self._full_path))
def __call__(
self,
env: EmbodiedEnv,
env_ids: torch.Tensor | Sequence[int] | slice | None,
entity_cfg: SceneEntityCfg | None = None,
folder_path: str | None = None,
*,
deployment: MaterializedAssetDeployment | None = None,
) -> None:
"""Apply the selected assets at an episode or batch boundary.
Args:
env: Existing environment whose selected objects are replaced.
env_ids: Reset rows; prepared deployments require all of B=1.
entity_cfg: Existing directory mode's mesh object selector.
folder_path: Existing directory mode's candidate directory/pattern.
deployment: Verified asset-batch artifact, or ``None`` while idle.
"""
if getattr(self, "_prepared_mode", False):
if entity_cfg is not None or folder_path is not None:
raise ValueError("Prepared replacement cannot also select a directory")
if deployment is None:
return
if isinstance(env_ids, slice):
selected_ids = list(range(env.num_envs))[env_ids]
elif env_ids is None:
selected_ids = list(range(env.num_envs))
else:
selection = torch.as_tensor(env_ids)
if selection.dtype not in (torch.int32, torch.int64):
raise ValueError(
"Prepared replacement env_ids must be integer rows"
)
selected_ids = selection.reshape(-1).tolist()
if env.num_envs != 1 or selected_ids != [0]:
raise ValueError(
"Prepared asset replacement requires a complete B=1 reset"
)
from embodichain.lab.gym.envs._asset_replacement import (
_apply_asset_deployment,
)
_apply_asset_deployment(env, deployment)
return
if deployment is not None:
raise ValueError(
"Directory replacement was not initialized for a deployment"
)
asset_path = random.choice(self._asset_group_path)
replacement_cfg = deepcopy(self.asset_cfg)
replacement_cfg.shape.fpath = asset_path
if self.asset_type == RigidObject:
env.sim.replace_rigid_object(cfg=replacement_cfg)
else:
logger.log_error("Only RigidObject assets are supported for replacement.")
[docs]
def register_entity_attrs(
env: EmbodiedEnv,
env_ids: torch.Tensor,
entity_cfg: SceneEntityCfg,
registration: str = "affordance_datas",
attrs: List[str] = [],
prefix: bool = True,
):
"""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.
Args:
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.
"""
entity = env.sim.get_asset(entity_cfg.uid)
if not hasattr(env, registration):
logger.log_warning(
f"Environment has no atrtribute {registration} for registration, please check again."
)
return
else:
registration_dict = getattr(env, registration, None)
if not isinstance(registration_dict, Dict):
logger.log_warning(
f"Got registration env.{registration} with type {type(registration_dict)}, please check again."
)
return
for attr in attrs:
attr_key = f"{entity_cfg.uid}_{attr}" if prefix else attr
if (attr_val := getattr(entity, attr_key, None)) is not None:
registration_dict.update({attr_key: attr_val})
elif (
attr_val := getattr(
env.event_manager.get_functor("prepare_extra_attr"), "extra_attrs", {}
)
.get(entity_cfg.uid, {})
.get(attr)
) is not None:
registration_dict.update({attr_key: attr_val})
else:
logger.log_warning(
f"Attr {attr} for entity {entity_cfg.uid} has neither been found in entity attrbutes nor prepare_extra_attrs functor, skipping.."
)
def register_entity_pose(
env: EmbodiedEnv,
env_ids: torch.Tensor | None,
entity_cfg: SceneEntityCfg,
registration: str = "affordance_datas",
compute_relative: bool | List | str = "all_robots",
compute_pose_object_to_arena: bool = True,
to_matrix: bool = True,
):
update_registration_dict = {}
if not hasattr(env, registration):
logger.log_warning(
f"Environment has no atrtribute {registration} for registration, please check again."
)
return
else:
registration_dict = getattr(env, registration, None)
if not isinstance(registration_dict, Dict):
logger.log_warning(
f"Got registration env.{registration} with type {type(registration_dict)}, please check again."
)
return
entity_pose_name, entity_pose = get_pose(
env, env_ids, entity_cfg, return_name=True, to_matrix=True
)
update_registration_dict.update({entity_pose_name: entity_pose})
if compute_relative:
# transform other entity's pose to entity frame
relative_poses = {}
if compute_relative == True:
entity_uids = (
env.sim.get_articulation_uid_list()
+ env.sim.get_rigid_object_uid_list()
+ env.sim.get_robot_uid_list()
)
elif isinstance(compute_relative, (str, list)):
entity_uids = resolve_uids(env, compute_relative)
else:
logger.log_warning(
f"Compute relative pose option with type {type(compute_relative)} is not supported, using empty list for skipping.."
)
entity_uids = []
for other_entity_uid in entity_uids:
if other_entity_uid != entity_cfg.uid:
# TODO: this is only for asset
other_entity_pose = env.sim.get_asset(other_entity_uid).get_local_pose(
to_matrix=True
)[env_ids, :]
relative_pose = torch.bmm(pose_inv(entity_pose), other_entity_pose)
relative_poses.update(
{
f"{other_entity_uid}_pose_{entity_pose_name.replace('_pose', '')}": relative_pose
}
)
update_registration_dict.update(relative_poses)
entity = env.sim.get_asset(entity_cfg.uid)
if isinstance(entity, RigidObject):
extra_attr_functor = env.event_manager.get_functor("prepare_extra_attr")
entity_extra_attrs = getattr(extra_attr_functor, "extra_attrs", {}).get(
entity_cfg.uid, {}
)
for (
entity_extra_attr_key,
entity_extra_attr_val,
) in entity_extra_attrs.items():
if entity_extra_attr_key.endswith("_pose_object"):
entity_extra_attr_val = torch.as_tensor(
entity_extra_attr_val, device=env.device
)
if entity_extra_attr_val.ndim < 3:
logger.log_info(
f"Got xyz_quat pose {entity_extra_attr_key}: {entity_extra_attr_val}, transforming it to matrix.",
color="green",
)
entity_extra_attr_val = xyz_quat_to_4x4_matrix(
entity_extra_attr_val
)
update_registration_dict.update(
{
entity_cfg.uid
+ "_"
+ (entity_extra_attr_key): entity_extra_attr_val
}
)
if compute_pose_object_to_arena:
# matmul broadcasts a static (1, 4, 4) pose over all envs; bmm doesn't.
pose_arena = torch.matmul(entity_pose, entity_extra_attr_val)
update_registration_dict.update(
{
entity_cfg.uid
+ "_"
+ (
entity_extra_attr_key.replace("_pose_object", "_pose")
): pose_arena
}
)
else:
logger.log_warning(
f"Now compute_pose_object_to_arena only support RigidObject type entity, skipping.."
)
if not to_matrix:
for key, val in update_registration_dict.items():
update_registration_dict[key] = trans_matrix_to_xyz_quat(val)
registration_dict = getattr(env, registration, None)
if not isinstance(registration_dict, Dict):
logger.log_warning(
f"Got registration env.{registration} with type {type(registration_dict)}, please check again."
)
return
registration_dict.update(update_registration_dict)
[docs]
def register_info_to_env(
env: EmbodiedEnv,
env_ids: torch.Tensor | None,
registry: List[Dict],
registration: str = "affordance_datas",
sim_update: bool = False,
) -> None:
"""Register entity information without advancing physics by default.
Args:
env: Environment receiving the registered information.
env_ids: Environment indices to read. None selects all environments.
registry: Entity information declarations.
registration: Name of the environment's destination mapping.
sim_update: Explicitly run the legacy 100-step settle before reading.
This requires all environments because physics advances the World.
"""
if env_ids is None:
env_ids = torch.arange(env.num_envs, device=env.device)
if sim_update:
selected = torch.as_tensor(env_ids, dtype=torch.long, device=env.device)
all_ids = torch.arange(env.num_envs, device=env.device)
if not torch.equal(selected.sort().values, all_ids):
raise ValueError(
"sim_update=True requires all environments because settling "
"advances the entire physics World."
)
logger.log_info(
"Calling env.sim.update(100) for after-physics-applied object attributes..",
color="green",
)
env.sim.update(step=100)
else:
env.sim.sync_render_state()
for entity_registry in registry:
entity_cfg = SceneEntityCfg(**entity_registry["entity_cfg"])
logger.log_info(f"Registering {entity_cfg.uid}..", color="green")
if (entity_attrs := entity_registry.get("attrs")) is not None:
prefix = entity_registry.get("prefix", True)
register_entity_attrs(
env, env_ids, entity_cfg, registration, entity_attrs, prefix
)
if (
pose_register_params := entity_registry.get("pose_register_params")
) is not None:
register_entity_pose(
env, env_ids, entity_cfg, registration, **pose_register_params
)
def resolve_uids(env: EmbodiedEnv, entity_uids: list[str] | str) -> list[str]:
if isinstance(entity_uids, str):
if entity_uids == "all_objects":
entity_uids = (
env.sim.get_rigid_object_uid_list()
+ env.sim.get_articulation_uid_list()
)
elif entity_uids == "all_robots":
entity_uids = env.sim.get_robot_uid_list()
elif entity_uids == "all_sensors":
entity_uids = env.sim.get_sensor_uid_list()
else:
# logger.log_warning(f"Entity uids {entity_uids} not supported in ['all_objects', 'all_robots', 'all_sensors'], wrapping it as a list..")
entity_uids = [entity_uids]
elif isinstance(entity_uids, (list, set, tuple)):
entity_uids = list(entity_uids)
else:
logger.log_error(
f"Entity uids {entity_uids} with type {type(entity_uids)} not supported in [List[str], str], please check again."
)
return entity_uids
def resolve_dict(env: EmbodiedEnv, entity_dict: Dict):
for entity_key in list(entity_dict.keys()):
entity_val = entity_dict.pop(entity_key)
entity_uids = resolve_uids(env, entity_key)
for entity_uid in entity_uids:
entity_dict.update({entity_uid: deepcopy(entity_val)})
return entity_dict
def get_pose(
env: EmbodiedEnv,
env_ids: torch.Tensor,
entity_cfg: SceneEntityCfg,
return_name: bool = True,
to_matrix: bool = True,
):
entity = env.sim.get_asset(entity_cfg.uid)
if isinstance(entity, RigidObject):
entity_pose = entity.get_local_pose(to_matrix=to_matrix)[env_ids, :]
entity_pose_register_name = entity_cfg.uid + "_pose"
elif isinstance(entity, Robot):
_, control_parts = resolve_matching_names(
entity_cfg.control_parts, list(entity.control_parts.keys())
)
if len(control_parts) != 1:
logger.log_warning(
"Only 1 control part can be assigned for computing the robot pose, please check again. Skipping"
)
return None
entity_cfg.control_parts = control_parts
control_part = control_parts[0]
# Select envs first, then joints: indexing with both at once pairs the
# two index lists element-wise instead of taking every joint per env.
control_part_qpos = entity.get_qpos()[env_ids][
:, entity.get_joint_ids(control_part)
]
entity_pose = entity.compute_fk(
control_part_qpos, name=control_part, to_matrix=to_matrix, env_ids=env_ids
) # NOTE: now compute_fk returns arena pose
entity_pose_register_name = control_part + "_pose"
else:
logger.log_warning(
f"Entity with tyope {type(entity)} is not supported, please check again."
)
return None
if return_name:
return entity_pose_register_name, entity_pose
else:
return entity_pose
[docs]
def drop_rigid_object_group_sequentially(
env: EmbodiedEnv,
env_ids: torch.Tensor | None,
entity_cfg: SceneEntityCfg,
drop_position: List[float] = [0.0, 0.0, 1.0],
position_range: Tuple[List[float], List[float]] = (
[-0.1, -0.1, 0.0],
[0.1, 0.1, 0.0],
),
physics_step: int = 2,
) -> None:
"""Drop rigid object group from a specified height sequentially in the environment.
Args:
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.
"""
obj_group: RigidObjectGroup = env.sim.get_rigid_object_group(entity_cfg.uid)
if obj_group is None:
logger.log_error(
f"RigidObjectGroup with UID '{entity_cfg.uid}' not found in the simulation."
)
num_instance = len(env_ids)
num_objects = obj_group.num_objects
range_low = torch.tensor(position_range[0], device=env.device)
range_high = torch.tensor(position_range[1], device=env.device)
drop_pos = (
torch.tensor(drop_position, device=env.device)
.unsqueeze_(0)
.repeat(num_instance, 1)
)
drop_pose = torch.zeros((num_instance, 7), device=env.device)
drop_pose[:, 6] = 1.0 # w component of xyzw quaternion
drop_pose[:, :3] = drop_pos
for i in range(num_objects):
random_offset = sample_uniform(
lower=range_low,
upper=range_high,
size=(num_instance, 3),
device=env.device,
)
drop_pose_i = drop_pose.unsqueeze(1)
drop_pose_i[:, 0, :3] = drop_pos + random_offset
obj_group.set_local_pose(pose=drop_pose_i, env_ids=env_ids, obj_ids=[i])
env.sim.update(step=physics_step)
[docs]
def set_detached_uids_for_env_reset(
env: EmbodiedEnv,
env_ids: torch.Tensor | None,
uids: list[str],
) -> None:
"""Set the UIDs of objects that are detached from automatic reset in the environment.
Args:
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.
"""
env.add_detached_uids_for_reset(uids=uids)
[docs]
def create_rigid_constraint(
env: EmbodiedEnv,
env_ids: torch.Tensor | None,
obj_a_cfg: SceneEntityCfg,
obj_b_cfg: SceneEntityCfg,
name: str,
local_frame_a: np.ndarray | None = None,
local_frame_b: np.ndarray | None = None,
) -> None:
"""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
:meth:`SimulationManager.create_rigid_constraint`.
Args:
env: The environment instance.
env_ids: Target environment indices. None -> all envs.
obj_a_cfg: SceneEntityCfg pointing at the first RigidObject.
obj_b_cfg: SceneEntityCfg pointing at the second RigidObject.
name: Base constraint name; per-arena names derived by the sim layer.
local_frame_a: Local joint frame on object A. None -> identity (object
A's origin). Accepts (4,4) or (N,4,4).
local_frame_b: 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.
"""
obj_a = env.sim.get_asset(obj_a_cfg.uid)
obj_b = env.sim.get_asset(obj_b_cfg.uid)
if not isinstance(obj_a, RigidObject) or not isinstance(obj_b, RigidObject):
logger.log_error(
f"Constraint '{name}' requires two RigidObjects, but got "
f"{type(obj_a).__name__} and {type(obj_b).__name__}."
)
env.sim.create_rigid_constraint(
cfg=RigidConstraintCfg(
name=name,
rigid_object_a_uid=obj_a_cfg.uid,
rigid_object_b_uid=obj_b_cfg.uid,
local_frame_a=local_frame_a,
local_frame_b=local_frame_b,
),
env_ids=env_ids,
)
[docs]
def remove_rigid_constraint(
env: EmbodiedEnv,
env_ids: torch.Tensor | None,
name: str,
) -> None:
"""Remove the named constraint for the given env_ids.
Delegates to :meth:`SimulationManager.remove_rigid_constraint`. Idempotent:
warns (via the sim layer) if the constraint is not found.
Args:
env: The environment instance.
env_ids: Target environment indices. None -> all envs.
name: Base constraint name to remove.
"""
env.sim.remove_rigid_constraint(name, env_ids=env_ids)