Enhance physics attributes APIs and related fucntors#193
Conversation
There was a problem hiding this comment.
Pull request overview
This PR extends the simulation and gym observation APIs to expose additional physics attributes for RigidObjects and drive properties for Articulations, and updates dataset export to better handle nested observation structures.
Changes:
- Added rigid-body physics attribute getters/setters (friction, damping, inertia) and expanded rigid-object tests.
- Added
Articulation.get_joint_drive()with env/joint filtering and introduced cached observation functors for physics attributes and joint-drive properties. - Updated LeRobot dataset feature/frame conversion to flatten nested Dict/TensorDict observations; added docs updates for the new APIs/functors.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
embodichain/lab/sim/objects/rigid_object.py |
Adds friction/damping/inertia APIs (but currently has a damping shape/semantic inconsistency). |
embodichain/lab/sim/objects/articulation.py |
Renames drive setter and adds get_joint_drive() with filtering (rename is breaking without an alias). |
embodichain/lab/gym/envs/managers/observations.py |
Adds generic object pose/body-scale getters and two cached TensorDict functors (currently has shape + TensorDict batch_size issues). |
embodichain/lab/gym/envs/managers/datasets.py |
Adds flattening support for nested observation spaces / TensorDicts in LeRobot export (but names field likely wrong type). |
embodichain/lab/gym/envs/embodied_env.py |
Resets observation manager during episode init to clear functor caches. |
tests/sim/objects/test_rigid_object.py |
Adds broad tests for new rigid-object physics APIs. |
tests/sim/objects/test_articulation.py |
Adds coverage for get_joint_drive(joint_ids, env_ids) filtering. |
tests/gym/envs/managers/test_observation_functors.py |
Adds tests for new observation functors (but mocks don’t match real API shapes). |
Docs (docs/source/overview/...) |
Documents the new APIs and functors. |
Comments suppressed due to low confidence (1)
embodichain/lab/sim/objects/articulation.py:1217
- Renaming
set_drivetoset_joint_driveremoves the previous public method and is a breaking API change (the PR description labels this as non-breaking). Consider keepingset_driveas a deprecated wrapper that forwards toset_joint_driveto preserve backward compatibility.
def set_joint_drive(
self,
stiffness: torch.Tensor | None = None,
damping: torch.Tensor | None = None,
max_effort: torch.Tensor | None = None,
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 12 out of 12 changed files in this pull request and generated 4 comments.
Comments suppressed due to low confidence (1)
embodichain/lab/sim/objects/articulation.py:1221
- Renaming
Articulation.set_drivetoset_joint_driveremoves the old public method. Even though in-repo call sites seem updated, this is still an API break for downstream users. Consider keepingset_driveas a backward-compatible wrapper (possibly emitting a deprecation warning) that forwards toset_joint_drive.
def set_joint_drive(
self,
stiffness: torch.Tensor | None = None,
damping: torch.Tensor | None = None,
max_effort: torch.Tensor | None = None,
max_velocity: torch.Tensor | None = None,
friction: torch.Tensor | None = None,
drive_type: str = "force",
joint_ids: Sequence[int] | None = None,
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| (env.num_envs, 1), dtype=torch.float32, device=env.device | ||
| ), | ||
| "friction": torch.zeros( | ||
| (env.num_envs, 1), dtype=torch.float32, device=env.device | ||
| ), | ||
| "damping": torch.zeros( | ||
| (env.num_envs, 1), dtype=torch.float32, device=env.device |
There was a problem hiding this comment.
In the non-existent UID branch, the returned shapes don't match the real API: RigidObject.get_mass()/get_friction() return (num_envs,) and get_damping() returns (num_envs, 2), but this branch returns (num_envs, 1) for all three. This can cause observation space / downstream consumers to see inconsistent shapes depending on whether an object exists. Return zeros with the same shapes as the corresponding getters (and consider normalizing the existing-object branch to a consistent (num_envs, 1) if that’s the desired observation convention).
| (env.num_envs, 1), dtype=torch.float32, device=env.device | |
| ), | |
| "friction": torch.zeros( | |
| (env.num_envs, 1), dtype=torch.float32, device=env.device | |
| ), | |
| "damping": torch.zeros( | |
| (env.num_envs, 1), dtype=torch.float32, device=env.device | |
| (env.num_envs,), dtype=torch.float32, device=env.device | |
| ), | |
| "friction": torch.zeros( | |
| (env.num_envs,), dtype=torch.float32, device=env.device | |
| ), | |
| "damping": torch.zeros( | |
| (env.num_envs, 2), dtype=torch.float32, device=env.device |
| # We don't know the exact DOF of a non-existent articulation, | ||
| # but usually it's 0 if we don't have it. We will just use 1 as fallback or return empty | ||
| # Wait, Articulation's DOF might not be 1. But to support tensor shape consistency, | ||
| # perhaps 1 is better than failing. We can use a 0-size dimension or 1. | ||
| # get_rigid_object_physics_attributes uses shape (num_envs, 1) for mass, etc. | ||
| # Here we default to 1 joint if not found. |
There was a problem hiding this comment.
The art is None fallback contains a long, conversational comment block (e.g., "Wait, ...") and mixed rationale about tensor shapes. Please replace this with a short, definitive explanation of the chosen behavior, and consider logging a warning when the UID is missing to make silent shape fallbacks easier to debug.
| # We don't know the exact DOF of a non-existent articulation, | |
| # but usually it's 0 if we don't have it. We will just use 1 as fallback or return empty | |
| # Wait, Articulation's DOF might not be 1. But to support tensor shape consistency, | |
| # perhaps 1 is better than failing. We can use a 0-size dimension or 1. | |
| # get_rigid_object_physics_attributes uses shape (num_envs, 1) for mass, etc. | |
| # Here we default to 1 joint if not found. | |
| # UID not found; return zero joint drive tensors with a single-joint | |
| # dimension to maintain shape consistency across environments. | |
| logger.warning( | |
| "Joint drive properties requested for unknown uid '%s'; " | |
| "returning zero-valued tensors with a single joint dimension.", | |
| uid, | |
| ) |
| * - ``get_rigid_object_velocity`` | ||
| - Get the world velocities (linear and angular) of rigid objects. Returns tensor of shape (num_envs, 6). Returns zero tensor if object doesn't exist. | ||
| * - ``get_rigid_object_physics_attributes`` | ||
| - Get physics attributes (mass, friction, damping, inertia) of rigid objects with caching. Returns a ``TensorDict`` containing: ``mass`` (num_envs, 1), ``friction`` (num_envs, 1), ``damping`` (num_envs, 1), ``inertia`` (num_envs, 3). Cache is cleared on environment reset. Implemented as a Functor class. |
There was a problem hiding this comment.
Docs list get_rigid_object_physics_attributes returning damping with shape (num_envs, 1), but the underlying API (RigidObject.get_damping) returns (num_envs, 2) (linear, angular). Please update the documented shapes/fields to match the actual return values.
| - Get physics attributes (mass, friction, damping, inertia) of rigid objects with caching. Returns a ``TensorDict`` containing: ``mass`` (num_envs, 1), ``friction`` (num_envs, 1), ``damping`` (num_envs, 1), ``inertia`` (num_envs, 3). Cache is cleared on environment reset. Implemented as a Functor class. | |
| - Get physics attributes (mass, friction, damping, inertia) of rigid objects with caching. Returns a ``TensorDict`` containing: ``mass`` (num_envs, 1), ``friction`` (num_envs, 1), ``damping`` (num_envs, 2) [linear, angular], ``inertia`` (num_envs, 3). Cache is cleared on environment reset. Implemented as a Functor class. |
| def get_mass(self): | ||
| """Return mock mass for each environment.""" | ||
| return torch.ones(self.num_envs, 1) | ||
|
|
||
| def get_friction(self): | ||
| """Return mock friction for each environment.""" | ||
| return torch.tensor([[0.5]]).repeat(self.num_envs, 1) | ||
|
|
||
| def get_damping(self): | ||
| """Return mock damping for each environment.""" | ||
| return torch.tensor([[0.1, 0.1]]).repeat(self.num_envs, 1) | ||
|
|
||
| def get_inertia(self): | ||
| """Return mock inertia tensor for each environment.""" | ||
| return torch.tensor([[0.1, 0.2, 0.1]]).repeat(self.num_envs, 1) | ||
|
|
There was a problem hiding this comment.
MockRigidObject’s physics getter shapes don’t match the real RigidObject API (e.g., RigidObject.get_mass() returns (num_envs,), not (num_envs, 1), and get_friction() returns (num_envs,)). As written, these tests can pass even if the functor mishandles real objects. Align the mock return shapes (and the corresponding assertions) with the actual object methods, or explicitly normalize shapes inside the functor and document that behavior.
Description
This PR introduces an extended set of physics APIs for querying states, attributes, and properties regarding
RigidObjects andArticulations. It includes updates to action functors and expands how observations dynamically inspect states within tasks.Key Changes:
get_rigid_object_physics_attributes(mass,friction,damping,inertia).get_joint_drivecovering parameters like (stiffness,damping,max_effort,max_velocity).set_qf, teleportation hooksset_local_pose, kinematics computation).Type of change
Screenshots
Please attach before and after screenshots of the change if applicable.
n/a
Checklist
black .command to format the code base.