Add joint armature support for articulations

.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"

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

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,

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]]

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,13 +1517,19 @@ 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,
                 damping_i,
                 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,
         )
 

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,
         )
 

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 = {}

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,13 +264,15 @@ 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]
         expected_damping = all_damping[env_ids][:, joint_ids]
         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."""