From 6878d7eaa2621479af114e24b93f52c1bfaae217 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 4 Jun 2026 11:47:25 +0800 Subject: [PATCH] feat(sim): add joint armature support for articulations Expose DexSim 0.4.1 articulation armature via JointDrivePropertiesCfg, Articulation/Robot drive APIs, and joint-drive observations. Extend get_joint_drive to return armature and update CI image for py311 DexSim. Co-authored-by: Cursor --- .github/workflows/main.yml | 2 +- docs/source/overview/sim/sim_articulation.md | 3 +- .../lab/gym/envs/managers/observations.py | 6 ++- embodichain/lab/sim/cfg.py | 9 ++++ embodichain/lab/sim/objects/articulation.py | 53 +++++++++++++++++-- embodichain/lab/sim/objects/robot.py | 5 ++ .../managers/test_observation_functors.py | 9 +++- tests/sim/objects/test_articulation.py | 16 ++++-- 8 files changed, 92 insertions(+), 11 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index f97b4fa2..3e21a3a9 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -16,7 +16,7 @@ jobs: NVIDIA_VISIBLE_DEVICES: all NVIDIA_DISABLE_REQUIRE: 1 container: &container_template - image: 192.168.3.13:5000/dexsdk:ubuntu22.04-cuda12.8.0-h5ffmpeg-v3 + image: 192.168.3.13:5000/dexsdk:ubuntu22.04-cuda12.8.0-h5ffmpeg-v3-py311 volumes: - "/cache:/cache" - "/usr/share/vulkan/icd.d:/usr/share/vulkan/icd.d" diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index f2edfc29..e8605047 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -65,6 +65,7 @@ The `drive_props` parameter controls the joint physics behavior. It is defined u | `max_effort` | `float` / `Dict` | `1.0e10` | Maximum effort (force/torque) the joint can exert. | | `max_velocity` | `float` / `Dict` | `1.0e10` | Maximum velocity allowed for the joint ($m/s$ or $rad/s$). | | `friction` | `float` / `Dict` | `0.0` | Joint friction coefficient. | +| `armature` | `float` / `Dict` | `0.0` | Joint armature added to joint-space inertia ($kg$ for prismatic, $kg \cdot m^2$ for revolute). | | `drive_type` | `str` | `"none"` | Drive mode: `"force"`(driven by a force), `"acceleration"`(driven by an acceleration) or `none`(no force). | ### Setup & Initialization @@ -138,7 +139,7 @@ State data is accessed via getter methods that return batched tensors (`N` envir | `get_link_pose(link_name, to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Specific link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | | `get_qpos(target=False)` | `(N, dof)` | Current joint positions (or joint targets if `target=True`). | | `get_qvel(target=False)` | `(N, dof)` | Current joint velocities (or velocity targets if `target=True`). | -| `get_joint_drive()` | `Tuple[Tensor, ...]` | Returns `(stiffness, damping, max_effort, max_velocity, friction)`, each shaped `(N, dof)`. | +| `get_joint_drive()` | `Tuple[Tensor, ...]` | Returns `(stiffness, damping, max_effort, max_velocity, friction, armature)`, each shaped `(N, dof)`. | ```python # Example: Accessing state diff --git a/embodichain/lab/gym/envs/managers/observations.py b/embodichain/lab/gym/envs/managers/observations.py index 50724cea..d3fcb984 100644 --- a/embodichain/lab/gym/envs/managers/observations.py +++ b/embodichain/lab/gym/envs/managers/observations.py @@ -1149,12 +1149,15 @@ def __call__( "friction": torch.zeros( (env.num_envs, 1), dtype=torch.float32, device=env.device ), + "armature": torch.zeros( + (env.num_envs, 1), dtype=torch.float32, device=env.device + ), }, batch_size=[env.num_envs], device=env.device, ) else: - stiffness, damping, max_effort, max_velocity, friction = ( + stiffness, damping, max_effort, max_velocity, friction, armature = ( art.get_joint_drive() ) result = TensorDict( @@ -1164,6 +1167,7 @@ def __call__( "max_effort": max_effort, "max_velocity": max_velocity, "friction": friction, + "armature": armature, }, batch_size=[env.num_envs], device=env.device, diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 4e4e0684..157c453a 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -693,6 +693,15 @@ class JointDrivePropertiesCfg: friction: Union[Dict[str, float], float] = 0.0 """Friction coefficient of the joint""" + armature: Union[Dict[str, float], float] = 0.0 + """Joint armature added to joint-space spatial inertia. + + Units depend on the joint model: + + * For prismatic (linear) joints, the unit is mass [kg]. + * For revolute (angular) joints, the unit is mass * scene_length^2 [kg-m^2]. + """ + @classmethod def from_dict( cls, init_dict: Dict[str, Union[str, float, int]] diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 15d377b7..513ee175 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -482,6 +482,19 @@ def joint_friction(self) -> torch.Tensor: device=self.device, ) + @property + def joint_armature(self) -> torch.Tensor: + """Get the joint armature of the articulation. + + Returns: + torch.Tensor: The joint armature of the articulation with shape (N, dof). + """ + return torch.as_tensor( + np.array([entity.get_drive()[5] for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + @cached_property def qpos_limits(self) -> torch.Tensor: """Get the joint position limits of the articulation. @@ -629,12 +642,19 @@ def __init__( dtype=torch.float32, device=device, ) + self.default_joint_armature = torch.full( + (num_entities, dof), + default_cfg.armature, + dtype=torch.float32, + device=device, + ) self._set_default_joint_drive() else: # Read current properties from USD-loaded entities self.default_joint_stiffness = self._data.joint_stiffness.clone() self.default_joint_damping = self._data.joint_damping.clone() self.default_joint_friction = self._data.joint_friction.clone() + self.default_joint_armature = self._data.joint_armature.clone() self.default_joint_max_effort = self._data.qf_limits.clone() self.default_joint_max_velocity = self._data.qvel_limits.clone() @@ -649,6 +669,9 @@ def __init__( usd_drive_pros.friction = ( self.default_joint_friction[0].cpu().numpy().tolist() ) + usd_drive_pros.armature = ( + self.default_joint_armature[0].cpu().numpy().tolist() + ) usd_drive_pros.max_effort = ( self.default_joint_max_effort[0].cpu().numpy().tolist() ) @@ -1391,6 +1414,7 @@ def set_joint_drive( max_effort: torch.Tensor | None = None, max_velocity: torch.Tensor | None = None, friction: torch.Tensor | None = None, + armature: torch.Tensor | None = None, drive_type: str = "none", joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None, @@ -1403,6 +1427,7 @@ def set_joint_drive( max_effort (torch.Tensor): The maximum effort of the joint drive with shape (len(env_ids), len(joint_ids)). max_velocity (torch.Tensor): The maximum velocity of the joint drive with shape (len(env_ids), len(joint_ids)). friction (torch.Tensor): The joint friction coefficient with shape (len(env_ids), len(joint_ids)). + armature (torch.Tensor): The joint armature with shape (len(env_ids), len(joint_ids)). drive_type (str, optional): The type of drive to apply. Defaults to "force". joint_ids (Sequence[int] | None, optional): The joint indices to apply the drive to. If None, applies to all joints. Defaults to None. env_ids (Sequence[int] | None, optional): The environment indices to apply the drive to. If None, applies to all environments. Defaults to None. @@ -1425,13 +1450,22 @@ def set_joint_drive( drive_args["max_velocity"] = max_velocity[i].cpu().numpy() if friction is not None: drive_args["joint_friction"] = friction[i].cpu().numpy() + if armature is not None: + drive_args["armature"] = armature[i].cpu().numpy() self._entities[env_idx].set_drive(**drive_args) def get_joint_drive( self, joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: """Get the drive properties for the articulation. Args: @@ -1441,8 +1475,8 @@ def get_joint_drive( If None, gets for all environments. Defaults to None. Returns: - Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing the stiffness, - damping, max_effort, max_velocity, and friction tensors with shape (N, len(joint_ids)) + Tuple[torch.Tensor, ...]: A tuple containing the stiffness, damping, max_effort, + max_velocity, friction, and armature tensors with shape (N, len(joint_ids)) for the specified environments. """ local_env_ids = self._all_indices if env_ids is None else env_ids @@ -1483,6 +1517,11 @@ def get_joint_drive( dtype=torch.float32, device=self.device, ) + armature = torch.zeros( + (len(local_env_ids), len(local_joint_ids)), + dtype=torch.float32, + device=self.device, + ) for i, env_idx in enumerate(local_env_ids): ( stiffness_i, @@ -1490,6 +1529,7 @@ def get_joint_drive( max_effort_i, max_velocity_i, friction_i, + armature_i, *_, ) = self._entities[env_idx].get_drive() stiffness[i] = torch.as_tensor( @@ -1507,7 +1547,10 @@ def get_joint_drive( friction[i] = torch.as_tensor( friction_i, dtype=torch.float32, device=self.device )[local_joint_ids_tensor] - return stiffness, damping, max_effort, max_velocity, friction + armature[i] = torch.as_tensor( + armature_i, dtype=torch.float32, device=self.device + )[local_joint_ids_tensor] + return stiffness, damping, max_effort, max_velocity, friction, armature def get_user_ids( self, link_name: str | None = None, env_ids: Sequence[int] | None = None @@ -1669,6 +1712,7 @@ def _set_default_joint_drive(self) -> None: ("max_effort", self.default_joint_max_effort), ("max_velocity", self.default_joint_max_velocity), ("friction", self.default_joint_friction), + ("armature", self.default_joint_armature), ] for prop_name, default_array in drive_props: @@ -1701,6 +1745,7 @@ def _set_default_joint_drive(self) -> None: max_effort=self.default_joint_max_effort, max_velocity=self.default_joint_max_velocity, friction=self.default_joint_friction, + armature=self.default_joint_armature, drive_type=drive_type, ) diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index 07273e80..e5e910a9 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -801,6 +801,7 @@ def set_joint_drive( max_effort: torch.Tensor | None = None, max_velocity: torch.Tensor | None = None, friction: torch.Tensor | None = None, + armature: torch.Tensor | None = None, drive_type: str = "force", joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None, @@ -814,6 +815,7 @@ def set_joint_drive( max_effort (torch.Tensor): The maximum effort of the joint drive with shape (len(env_ids), len(joint_ids)). max_velocity (torch.Tensor): The maximum velocity of the joint drive with shape (len(env_ids), len(joint_ids)). friction (torch.Tensor): The joint friction coefficient with shape (len(env_ids), len(joint_ids)). + armature (torch.Tensor): The joint armature with shape (len(env_ids), len(joint_ids)). drive_type (str, optional): The type of drive to apply. Defaults to "force". joint_ids (Sequence[int] | None, optional): The joint indices to apply the drive to. If None, applies to all joints. Defaults to None. env_ids (Sequence[int] | None, optional): The environment indices to apply the drive to. If None, applies to all environments. Defaults to None. @@ -824,6 +826,7 @@ def set_joint_drive( max_effort=max_effort, max_velocity=max_velocity, friction=friction, + armature=armature, drive_type=drive_type, joint_ids=joint_ids, env_ids=env_ids, @@ -840,6 +843,7 @@ def _set_default_joint_drive(self) -> None: ("max_effort", self.default_joint_max_effort), ("max_velocity", self.default_joint_max_velocity), ("friction", self.default_joint_friction), + ("armature", self.default_joint_armature), ] for prop_name, default_array in drive_props: @@ -894,6 +898,7 @@ def _set_default_joint_drive(self) -> None: max_effort=self.default_joint_max_effort, max_velocity=self.default_joint_max_velocity, friction=self.default_joint_friction, + armature=self.default_joint_armature, drive_type=drive_type, ) diff --git a/tests/gym/envs/managers/test_observation_functors.py b/tests/gym/envs/managers/test_observation_functors.py index a9238d90..ced6e1f7 100644 --- a/tests/gym/envs/managers/test_observation_functors.py +++ b/tests/gym/envs/managers/test_observation_functors.py @@ -77,7 +77,8 @@ def get_joint_drive(self, joint_ids=None, env_ids=None): max_effort = torch.ones((num_envs, joints), device=self.device) * 50.0 max_velocity = torch.ones((num_envs, joints), device=self.device) * 5.0 friction = torch.ones((num_envs, joints), device=self.device) * 1.0 - return stiffness, damping, max_effort, max_velocity, friction + armature = torch.ones((num_envs, joints), device=self.device) * 0.5 + return stiffness, damping, max_effort, max_velocity, friction, armature class MockRigidObject: @@ -673,12 +674,14 @@ def test_returns_correct_shapes(self): assert "max_effort" in result.keys() assert "max_velocity" in result.keys() assert "friction" in result.keys() + assert "armature" in result.keys() assert result["stiffness"].shape == (4, 6) assert result["damping"].shape == (4, 6) assert result["max_effort"].shape == (4, 6) assert result["max_velocity"].shape == (4, 6) assert result["friction"].shape == (4, 6) + assert result["armature"].shape == (4, 6) def test_returns_correct_values(self): """Test that the functor returns expected mock values.""" @@ -695,6 +698,7 @@ def test_returns_correct_values(self): assert torch.allclose(result["max_effort"], torch.ones(4, 6) * 50.0) assert torch.allclose(result["max_velocity"], torch.ones(4, 6) * 5.0) assert torch.allclose(result["friction"], torch.ones(4, 6) * 1.0) + assert torch.allclose(result["armature"], torch.ones(4, 6) * 0.5) def test_returns_zeros_for_nonexistent_object(self): """Test that zeros are returned for non-existent objects.""" @@ -711,6 +715,7 @@ def test_returns_zeros_for_nonexistent_object(self): assert torch.allclose(result["max_effort"], torch.zeros(4, 1)) assert torch.allclose(result["max_velocity"], torch.zeros(4, 1)) assert torch.allclose(result["friction"], torch.zeros(4, 1)) + assert torch.allclose(result["armature"], torch.zeros(4, 1)) def test_caches_data_across_calls(self): """Test that fetched data is cached for subsequent calls.""" @@ -723,6 +728,7 @@ def test_caches_data_across_calls(self): torch.ones(4, 6), torch.ones(4, 6), torch.ones(4, 6), + torch.ones(4, 6), ) ) obs = {} @@ -748,6 +754,7 @@ def test_reset_clears_cache(self): torch.ones(4, 6), torch.ones(4, 6), torch.ones(4, 6), + torch.ones(4, 6), ) ) obs = {} diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 8f9b42a1..89c37e72 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -237,9 +237,14 @@ def test_setter_methods(self): def test_get_joint_drive_with_joint_ids(self): """Test get_joint_drive supports joint_ids and env_ids filtering.""" - all_stiffness, all_damping, all_max_effort, all_max_velocity, all_friction = ( - self.art.get_joint_drive() - ) + ( + all_stiffness, + all_damping, + all_max_effort, + all_max_velocity, + all_friction, + all_armature, + ) = self.art.get_joint_drive() assert all_stiffness.shape == ( NUM_ARENAS, @@ -259,6 +264,7 @@ def test_get_joint_drive_with_joint_ids(self): max_effort, max_velocity, friction, + armature, ) = self.art.get_joint_drive(joint_ids=joint_ids, env_ids=env_ids) expected_stiffness = all_stiffness[env_ids][:, joint_ids] @@ -266,6 +272,7 @@ def test_get_joint_drive_with_joint_ids(self): expected_max_effort = all_max_effort[env_ids][:, joint_ids] expected_max_velocity = all_max_velocity[env_ids][:, joint_ids] expected_friction = all_friction[env_ids][:, joint_ids] + expected_armature = all_armature[env_ids][:, joint_ids] expected_shape = (len(env_ids), len(joint_ids)) assert ( @@ -286,6 +293,9 @@ def test_get_joint_drive_with_joint_ids(self): assert torch.allclose( friction, expected_friction, atol=1e-5 ), "FAIL: friction does not match expected filtered values" + assert torch.allclose( + armature, expected_armature, atol=1e-5 + ), "FAIL: armature does not match expected filtered values" def teardown_method(self): """Clean up resources after each test method."""