# ----------------------------------------------------------------------------
# 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 torch
import dexsim
import numpy as np
from functools import cached_property
from dataclasses import dataclass
from typing import List, Sequence, Union
from dexsim.models import MeshObject
from dexsim.engine import ClothBody, PhysicsScene
from dexsim.types import ClothBodyGPUAPIReadWriteType
from embodichain.lab.sim.common import (
BatchEntity,
)
from embodichain.lab.sim.material import (
VisualMaterial,
VisualMaterialInst,
_capture_render_materials,
_restore_render_materials,
_wrap_first_render_material,
)
from embodichain.utils.math import (
matrix_from_euler,
)
from embodichain.utils import logger
from embodichain.lab.sim.cfg import (
ClothObjectCfg,
)
from embodichain.utils.math import xyz_quat_to_4x4_matrix
[docs]
@dataclass
class ClothBodyData:
"""Data manager for cloth.
Note:
1. The pose data managed by dexsim is in the format of (qx, qy, qz, qw, x, y, z), but in EmbodiChain, we use (x, y, z, qw, qx, qy, qz) format.
"""
[docs]
def __init__(
self, entities: List[MeshObject], ps: PhysicsScene, device: torch.device
) -> None:
"""Initialize the ClothBodyData.
Args:
entities (List[MeshObject]): List of MeshObjects representing the cloth bodies.
ps (PhysicsScene): The physics scene.
device (torch.device): The device to use for the cloth body data.
"""
self.entities = entities
# TODO: cloth body data can only be stored in cuda device for now.
self.device = device
# TODO: inorder to retrieve arena position, we need to access the node of each entity.
self.ps = ps
self.num_instances = len(entities)
self.cloth_bodies: Sequence[ClothBody] = [
self.entities[i].get_physical_body() for i in range(self.num_instances)
]
self.n_vertices = self.cloth_bodies[0].get_num_vertices()
self._rest_position_buffer = torch.empty(
(self.num_instances, self.n_vertices, 4),
device=self.device,
dtype=torch.float32,
)
for i, cloth_body in enumerate(self.cloth_bodies):
self._rest_position_buffer[i] = cloth_body.get_position_inv_mass_buffer()
self._vertex_position = torch.zeros(
(self.num_instances, self.n_vertices, 3),
device=self.device,
dtype=torch.float32,
)
self._vertex_velocity = torch.zeros(
(self.num_instances, self.n_vertices, 3),
device=self.device,
dtype=torch.float32,
)
@property
def rest_vertices(self):
"""Get the rest position buffer of the cloth bodies."""
return self._rest_position_buffer[:, :, :3].clone()
@property
def vertex_position(self):
"""Get the current vertex position buffer of the cloth bodies."""
for i, clothbody in enumerate(self.cloth_bodies):
self._vertex_position[i] = clothbody.get_position_inv_mass_buffer()[:, :3]
return self._vertex_position.clone()
@property
def vertex_velocity(self):
"""Get the current vertex velocity buffer of the cloth bodies."""
for i, clothbody in enumerate(self.cloth_bodies):
self._vertex_velocity[i] = clothbody.get_velocity_buffer()[:, 3:]
return self._vertex_velocity.clone()
[docs]
class ClothObject(BatchEntity):
"""ClothObject represents a batch of cloth body in the simulation."""
[docs]
def __init__(
self,
cfg: ClothObjectCfg,
entities: List[MeshObject] = None,
device: torch.device = torch.device("cpu"),
) -> None:
self._world: dexsim.World = dexsim.default_world()
self._ps = self._world.get_physics_scene()
self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist()
self._data = ClothBodyData(entities=entities, ps=self._ps, device=device)
self._world.update(0.001)
self._visual_material: List[VisualMaterialInst | None] = [None] * len(entities)
self.is_shared_visual_material = False
super().__init__(cfg=cfg, entities=entities, device=device)
self._initialize_existing_visual_material()
self._set_default_collision_filter()
def _initialize_existing_visual_material(self) -> None:
"""Wrap asset-parsed materials during cloth-object construction.
For a multi-segment render body, the first segment with a valid
material is registered as the environment's representative material.
"""
self._original_visual_material = [[] for _ in self._entities]
self._original_visual_material_inst = [None] * len(self._entities)
for env_idx, entity in enumerate(self._entities):
render_body = entity.get_render_body()
if render_body is None:
continue
original_materials = _capture_render_materials(render_body)
self._original_visual_material[env_idx] = original_materials
wrapped = _wrap_first_render_material(original_materials)
if wrapped is not None:
self._visual_material[env_idx] = wrapped
self._original_visual_material_inst[env_idx] = wrapped
[docs]
def set_visual_material(
self,
mat: VisualMaterial,
env_ids: Sequence[int] | None = None,
shared: bool = False,
) -> None:
"""Set visual material for the cloth object.
Args:
mat: The material template to assign.
env_ids: Environment indices. If None, all instances are used.
shared: Whether selected environments share one material instance.
"""
local_env_ids = self._all_indices if env_ids is None else env_ids
if shared:
if len(local_env_ids) != self.num_instances:
logger.log_error("Cannot share material instance for partial env_ids.")
mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}")
for env_idx in local_env_ids:
self._entities[env_idx].set_material(mat_inst.mat)
self._visual_material[env_idx] = mat_inst
self.is_shared_visual_material = True
else:
for env_idx in local_env_ids:
mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{env_idx}")
self._entities[env_idx].set_material(mat_inst.mat)
self._visual_material[env_idx] = mat_inst
self.is_shared_visual_material = False
[docs]
def restore_visual_material(self, env_ids: Sequence[int] | None = None) -> None:
"""Restore visual materials captured when the cloth object was created.
Args:
env_ids: Environment indices. If None, all instances are restored.
"""
if not hasattr(self, "_original_visual_material"):
return
local_env_ids = self._all_indices if env_ids is None else env_ids
for env_idx in local_env_ids:
render_body = self._entities[env_idx].get_render_body()
if render_body is None:
continue
_restore_render_materials(
render_body, self._original_visual_material[env_idx]
)
self._visual_material[env_idx] = self._original_visual_material_inst[
env_idx
]
self.is_shared_visual_material = False
[docs]
def get_visual_material_inst(
self, env_ids: Sequence[int] | None = None
) -> List[VisualMaterialInst | None]:
"""Get the material instance registered for each selected environment.
Args:
env_ids: Environment indices. If None, all instances are returned.
Returns:
The existing material wrappers, or None where an asset has no material.
"""
ids = env_ids if env_ids is not None else range(self.num_instances)
return [self._visual_material[i] for i in ids]
def _set_default_collision_filter(self) -> None:
collision_filter_data = torch.zeros(
size=(self.num_instances, 4), dtype=torch.int32
)
for i in range(self.num_instances):
collision_filter_data[i, 0] = i
collision_filter_data[i, 1] = 1
self.set_collision_filter(collision_filter_data)
[docs]
def set_collision_filter(
self, filter_data: torch.Tensor, env_ids: Sequence[int] | None = None
) -> None:
"""Set collision filter data for the cloth object.
Args:
filter_data (torch.Tensor): [N, 4] of int.
First element of each object is arena id.
If 2nd element is 0, the object will collision with all other objects in world.
3rd and 4th elements are not used currently.
env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used.
"""
local_env_ids = self._all_indices if env_ids is None else env_ids
if len(local_env_ids) != len(filter_data):
logger.log_error(
f"Length of env_ids {len(local_env_ids)} does not match pose length {len(filter_data)}."
)
filter_data_np = filter_data.cpu().numpy().astype(np.uint32)
for i, env_idx in enumerate(local_env_ids):
self._entities[env_idx].get_physical_body().set_collision_filter_data(
filter_data_np[i]
)
@property
def body_data(self) -> ClothBodyData | None:
"""Get the cloth body data manager for this cloth object.
Returns:
ClothBodyData | None: The cloth body data manager.
"""
return self._data
[docs]
def get_rest_vertex_position(self) -> torch.Tensor:
"""Get the rest vertex position of the cloth bodies.
Returns:
torch.Tensor: The rest vertex position of the cloth bodies, shape (num_instances, n_vertices, 3).
"""
return self._data.rest_vertices
[docs]
def get_current_vertex_position(self) -> torch.Tensor:
"""Get the current vertex position of the cloth bodies.
Returns:
torch.Tensor: The current vertex position of the cloth bodies, shape (num_instances, n_vertices, 3).
"""
return self._data.vertex_position
[docs]
def get_current_vertex_velocity(self) -> torch.Tensor:
"""Get the current vertex velocity of the cloth bodies.
Returns:
torch.Tensor: The current vertex velocity of the cloth bodies, shape (num_instances, n_vertices, 3).
"""
return self._data.vertex_velocity
[docs]
def set_local_pose(
self, pose: torch.Tensor, env_ids: Sequence[int] | None = None
) -> None:
"""Set local pose of the cloth object.
Args:
pose (torch.Tensor): The local pose of the cloth object with shape (N, 7) or (N, 4, 4).
env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used.
"""
from embodichain.lab.sim import SimulationManager
sim = SimulationManager.get_instance()
local_env_ids = self._all_indices if env_ids is None else env_ids
if len(local_env_ids) != len(pose):
logger.log_error(
f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}."
)
if pose.dim() == 2 and pose.shape[1] == 7:
pose4x4 = xyz_quat_to_4x4_matrix(pose)
elif pose.dim() == 3 and pose.shape[1:3] == (4, 4):
pose4x4 = pose
else:
logger.log_error(
f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)."
)
arena_offsets = sim.arena_offsets
for i, env_idx in enumerate(local_env_ids):
# TODO: cloth body cannot directly set by `set_local_pose` currently.
rest_vertices = self.body_data.rest_vertices[i]
rotation = pose4x4[i][:3, :3]
translation = pose4x4[i][:3, 3]
# apply transformation to local rest vertices and back
rest_vertices_local = rest_vertices - arena_offsets[i]
transformed_vertices = rest_vertices_local @ rotation.T + translation
transformed_vertices = transformed_vertices + arena_offsets[i]
cloth_body: ClothBody = self._entities[env_idx].get_physical_body()
position_buffer = cloth_body.get_position_inv_mass_buffer()
velocity_buffer = cloth_body.get_velocity_buffer()
position_buffer[:, :3] = transformed_vertices
velocity_buffer[:, 3:] = 0.0
cloth_body.mark_dirty(ClothBodyGPUAPIReadWriteType.ALL)
# TODO: currently cloth body has no wake up interface, use set_wake_counter and pass in a positive value to wake it up
cloth_body.set_wake_counter(0.4)
[docs]
def get_local_pose(self, to_matrix=False):
"""Get local pose of the cloth object.
Args:
to_matrix (bool, optional): If True, return the pose as a 4x4 matrix. If False, return as (x, y, z, qw, qx, qy, qz). Defaults to False.
Returns:
torch.Tensor: The local pose of the cloth object with shape (N, 7) or (N, 4, 4) depending on `to_matrix`.
"""
raise NotImplementedError(
"Getting local pose for ClothObject is not supported."
)
[docs]
def reset(self, env_ids: Sequence[int] | None = None) -> None:
local_env_ids = self._all_indices if env_ids is None else env_ids
num_instances = len(local_env_ids)
self.restore_visual_material(env_ids=local_env_ids)
# TODO: set attr for cloth body after loading in physics scene.
# rest cloth body to init_pos
pos = torch.as_tensor(
self.cfg.init_pos, dtype=torch.float32, device=self.device
)
rot = (
torch.as_tensor(self.cfg.init_rot, dtype=torch.float32, device=self.device)
* torch.pi
/ 180.0
)
pos = pos.unsqueeze(0).repeat(num_instances, 1)
rot = rot.unsqueeze(0).repeat(num_instances, 1)
mat = matrix_from_euler(rot, "XYZ")
pose = (
torch.eye(4, dtype=torch.float32, device=self.device)
.unsqueeze(0)
.repeat(num_instances, 1, 1)
)
pose[:, :3, 3] = pos
pose[:, :3, :3] = mat
self.set_local_pose(pose, env_ids=local_env_ids)
[docs]
def destroy(self) -> None:
# TODO: not tested yet
env = self._world.get_env()
arenas = env.get_all_arenas()
if len(arenas) == 0:
arenas = [env]
for i, entity in enumerate(self._entities):
arenas[i].remove_actor(entity)