[FEATURE] Add public API to RigidEntity for kinematic and potential energy. by Lidang-Jiang · Pull Request #2613 · Genesis-Embodied-AI/Genesis

Lidang-Jiang · 2026-03-28T14:29:45Z

Summary

Add get_kinetic_energy(), get_potential_energy(), and get_total_energy() methods to RigidEntity
Kinetic energy uses generalized-coordinate mass matrix (KE = 0.5 * dq^T * M(q) * dq) for accurate computation on articulated bodies
Potential energy sums m_i * g^T * p_i over all links, supporting arbitrary gravity directions
All methods support envs_idx for parallel environments

API Demo Output

Full output

[Genesis] [22:27:00] [INFO] ╭───────────────────────────────────────────────╮
[Genesis] [22:27:00] [INFO] │┈┉┈┉┈┉┈┉┈┉┈┉┈┉┈┉┈┉┈ Genesis ┈┉┈┉┈┉┈┉┈┉┈┉┈┉┈┉┈┉┈│
[Genesis] [22:27:00] [INFO] ╰───────────────────────────────────────────────╯
[Genesis] [22:27:00] [INFO] Running on [11th Gen Intel(R) Core(TM) i7-11800H @ 2.30GHz] with backend gs.cpu. Device memory: 11.68 GB.
[Genesis] [22:27:00] [INFO] Genesis initialized. version: 0.4.3, theme: dark, seed: None, debug: False, precision: 64, performance: False, verbose: INFO
[Genesis] [22:27:02] [INFO] Scene <29f4fc1> created.
[Genesis] [22:27:02] [INFO] Adding <gs.engine.entities.RigidEntity>. idx: 0, uid: <522a7d0>, morph: <gs.morphs.Plane>, material: <gs.materials.Rigid>.
[Genesis] [22:27:08] [INFO] Adding <gs.engine.entities.RigidEntity>. idx: 1, uid: <4f36b3c>, morph: <gs.morphs.Box>, material: <gs.materials.Rigid>.
[Genesis] [22:27:08] [INFO] Building scene <29f4fc1>...
[Genesis] [22:27:09] [INFO] Compiling simulation kernels...
[Genesis] [22:27:10] [INFO] Building visualizer...
=== RigidEntity Energy API Demo ===
Initial position: tensor([0., 0., 1.])
Initial velocity: tensor([0., 0., 0.])
Kinetic Energy: 0.0
Potential Energy: 5.886000000000001
Total Energy: 5.886000000000001

After 50 steps free fall:
Position: tensor([0.0000, 0.0000, 0.9875])
Velocity: tensor([ 0.0000,  0.0000, -0.4905])
Kinetic Energy: 0.07217707499999997
Potential Energy: 5.812379383500009
Total Energy: 5.8845564585000085

Energy conservation check:
  Initial total energy: 5.886000
  Final total energy:   5.884556
  Relative drift:       0.024525%

Unit Test Results

pytest output (13 tests, all passed)

============================= test session starts ==============================
platform linux -- Python 3.11.14, pytest-9.0.2, pluggy-1.6.0 -- /home/devuser/workspace/public-oss/embodied-robotics/myenv-genesis-2309/bin/python
cachedir: .pytest_cache
rootdir: /home/devuser/workspace/public-oss/embodied-robotics/worktree-genesis-2309
configfile: pyproject.toml
plugins: xdist-3.8.0, print-1.2.2, anyio-4.13.0, syrupy-5.1.0, forked-1.6.0
created: 1/1 worker
1 worker [13 items]

scheduling tests via WorkStealingScheduling

tests/test_energy.py::TestEnergyFreeFall::test_stationary_object_has_zero_kinetic_energy
[gw0] [  7%] PASSED tests/test_energy.py::TestEnergyFreeFall::test_stationary_object_has_zero_kinetic_energy
tests/test_energy.py::TestEnergyFreeFall::test_potential_energy_proportional_to_height
[gw0] [ 15%] PASSED tests/test_energy.py::TestEnergyFreeFall::test_potential_energy_proportional_to_height
tests/test_energy.py::TestEnergyFreeFall::test_total_energy_equals_sum
[gw0] [ 23%] PASSED tests/test_energy.py::TestEnergyFreeFall::test_total_energy_equals_sum
tests/test_energy.py::TestEnergyFreeFall::test_energy_conservation_free_fall
[gw0] [ 30%] PASSED tests/test_energy.py::TestEnergyFreeFall::test_energy_conservation_free_fall
tests/test_energy.py::TestEnergyFreeFall::test_kinetic_energy_increases_during_free_fall
[gw0] [ 38%] PASSED tests/test_energy.py::TestEnergyFreeFall::test_kinetic_energy_increases_during_free_fall
tests/test_energy.py::TestEnergyFreeFall::test_potential_energy_decreases_during_free_fall
[gw0] [ 46%] PASSED tests/test_energy.py::TestEnergyFreeFall::test_potential_energy_decreases_during_free_fall
tests/test_energy.py::TestEnergyZeroGravity::test_zero_gravity_zero_potential
[gw0] [ 53%] PASSED tests/test_energy.py::TestEnergyZeroGravity::test_zero_gravity_zero_potential
tests/test_energy.py::TestEnergyZeroGravity::test_zero_gravity_kinetic_energy_conservation
[gw0] [ 61%] PASSED tests/test_energy.py::TestEnergyZeroGravity::test_zero_gravity_kinetic_energy_conservation
tests/test_energy.py::TestEnergyMultiLink::test_multi_link_energy_conservation
[gw0] [ 69%] PASSED tests/test_energy.py::TestEnergyMultiLink::test_multi_link_energy_conservation
tests/test_energy.py::TestEnergyReturnShape::test_single_env_returns_scalar
[gw0] [ 76%] PASSED tests/test_energy.py::TestEnergyReturnShape::test_single_env_returns_scalar
tests/test_energy.py::TestEnergyReturnShape::test_parallel_envs_returns_vector
[gw0] [ 84%] PASSED tests/test_energy.py::TestEnergyReturnShape::test_parallel_envs_returns_vector
tests/test_energy.py::TestEnergyReturnShape::test_parallel_envs_with_envs_idx
[gw0] [ 92%] PASSED tests/test_energy.py::TestEnergyReturnShape::test_parallel_envs_with_envs_idx
tests/test_energy.py::TestEnergyNonStandardGravity::test_gravity_along_x
[gw0] [100%] PASSED tests/test_energy.py::TestEnergyNonStandardGravity::test_gravity_along_x

============================== slowest durations ===============================

(39 durations < 100s hidden.)
======================== 13 passed in 82.27s (0:01:22) ========================

Test Plan

Free-fall energy conservation (drift < 0.005%)
Zero-gravity kinetic energy conservation (exact)
Multi-link articulated body energy conservation
Return shape validation (scalar / batched / envs_idx subset)
Non-standard gravity direction
Stationary object has zero kinetic energy
PE proportional to height
Total energy = KE + PE
KE increases / PE decreases during free fall

Add get_kinetic_energy(), get_potential_energy(), and get_total_energy() methods to RigidEntity for computing mechanical energy in simulation. - Kinetic energy uses generalized-coordinate mass matrix (handles all inertial coupling terms for articulated bodies) - Potential energy sums m_i * g^T * p_i over all links - All methods support envs_idx for parallel environments Closes Genesis-Embodied-AI#2309 Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>

duburcqa

This is invalid. It does not take into account motor armature.

- Add RigidSolver.recompute_mass_matrix() to get clean mass matrix (structural inertia + armature, no implicit damping) - get_kinetic_energy() now recomputes mass matrix before use - Delete test_energy.py; add 2 function-based tests to test_rigid_physics.py - Tests: no class, no precision requirement, marked @pytest.mark.required Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>

Lidang-Jiang · 2026-03-28T15:28:41Z

Thanks for the review! Addressed all feedback.

Motor armature fix:

The root issue was that get_mass_mat() returns the solver's shared mass matrix buffer, which after scene.step() with the default approximate_implicitfast integrator gets contaminated with implicit damping correction terms (M[i,i] += damping * dt + kv * dt, see forward_dynamics.py:527-541). This made the kinetic energy calculation incorrect.

Fix: Added RigidSolver.recompute_mass_matrix() which calls kernel_compute_mass_matrix with implicit_damping=False — this recomputes the mass matrix with structural inertia + motor armature on the diagonal, but without the implicit damping terms. get_kinetic_energy() now calls this before reading the mass matrix.

Test changes:

Deleted tests/test_energy.py (class-based, 16 tests, 64-bit only)
Added 2 function-based tests in tests/test_rigid_physics.py, marked @pytest.mark.required, no precision requirement

- Add unit [J] to all energy method titles - Remove redundant return descriptions - Remove invalid Note about mid-step calls - Add Note about mass matrix recomputation performance - Use 120-char linewidth for docstrings - Remove implementation detail about implicit damping Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>

- Merge test_energy_conservation_free_fall and test_energy_kinetic_potential_relation into one test - Add 2 spheres: undamped (dampratio=0) and damped (default) - Verify analytical KE/PE during free fall (semi-implicit Euler formulas) - Verify energy conservation for undamped sphere, strict decrease for damped - Track first ground impact timestep Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>

Require changes.

- Wrap docstrings to 120 chars linewidth - Skip mass matrix recomputation for non-approximate_implicitfast integrators - Remove recompute_mass_matrix() from RigidSolver, call kernel directly in RigidEntity - Parameterize test with both Euler and approximate_implicitfast integrators - One line per option in test entity creation - Use non-zero timeconst in undamped sol_params to avoid NaN on field backend Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>

duburcqa · 2026-04-02T13:24:33Z

@Lidang-Jiang Some comments are still open.

…mpact step

duburcqa · 2026-04-03T07:29:47Z

I will help you to investigate the nan issue.

github-actions · 2026-04-03T08:06:39Z

⚠️ Abnormal Benchmark Result Detected ➡️ Report

Lidang-Jiang · 2026-04-03T09:23:36Z

I will help you to investigate the nan issue.

thanks

…s near zero. Add eps parameter to imp_aref to prevent NaN from 1/(dampratio²) when dampratio ≈ 0. Reject dampratio < EPS in _sanitize_sol_params with a clear error message explaining that dampratio controls spring stiffness, not velocity damping. Document the impedance model semantics in default_solver_params docstring. Update energy conservation test to use nearly elastic contact params (timeconst=10, dampratio=0.001) with tighter tolerance.

…nergy. (Genesis-Embodied-AI#2613) Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com> Co-authored-by: Alexis Duburcq <alexis.duburcq@gmail.com>

duburcqa requested changes Mar 28, 2026

View reviewed changes

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread tests/test_energy.py Outdated

Lidang-Jiang requested a review from YilingQiao as a code owner March 28, 2026 15:20

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 28, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread tests/test_rigid_physics.py

duburcqa previously approved these changes Mar 29, 2026

View reviewed changes

duburcqa changed the title ~~[Feature] Add energy computation methods to RigidEntity~~ [Feature] Add public API to RigidEntity for computing kinematic and potential energy. Mar 29, 2026

duburcqa changed the title ~~[Feature] Add public API to RigidEntity for computing kinematic and potential energy.~~ [Feature] Add public API to RigidEntity for kinematic and potential energy. Mar 29, 2026

duburcqa changed the title ~~[Feature] Add public API to RigidEntity for kinematic and potential energy.~~ [FEATURE] Add public API to RigidEntity for kinematic and potential energy. Mar 29, 2026

duburcqa enabled auto-merge (squash) March 29, 2026 13:11

duburcqa disabled auto-merge March 29, 2026 13:13

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread genesis/engine/entities/rigid_entity/rigid_entity.py Outdated

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread genesis/engine/solvers/rigid/rigid_solver.py Outdated

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread tests/test_rigid_physics.py Outdated

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread tests/test_rigid_physics.py Outdated

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread tests/test_rigid_physics.py Outdated

duburcqa reviewed Mar 29, 2026

View reviewed changes

Comment thread tests/test_rigid_physics.py Outdated

duburcqa self-requested a review March 29, 2026 13:29

Lidang-Jiang added 2 commits April 2, 2026 22:01

refactor: add FIXME comment for timeconst=0 NaN issue on field backend

7f76609

refactor: remove .item() calls and use real collision detection for i…

4aae078

…mpact step

duburcqa force-pushed the feat/rigid-energy-api branch from 1f9c7cc to bcca404 Compare April 3, 2026 09:08

duburcqa approved these changes Apr 3, 2026

View reviewed changes

duburcqa previously approved these changes Apr 3, 2026

View reviewed changes

duburcqa dismissed their stale review via 7eeeaaf April 3, 2026 09:45

duburcqa force-pushed the feat/rigid-energy-api branch from bcca404 to 7eeeaaf Compare April 3, 2026 09:45

duburcqa force-pushed the feat/rigid-energy-api branch from 7eeeaaf to 4cc40d5 Compare April 3, 2026 09:52

duburcqa merged commit 725eb36 into Genesis-Embodied-AI:main Apr 3, 2026
19 of 20 checks passed

Conversation

Lidang-Jiang commented Mar 28, 2026

Summary

API Demo Output

Unit Test Results

Test Plan

Uh oh!

duburcqa left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Lidang-Jiang commented Mar 28, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

duburcqa commented Apr 2, 2026

Uh oh!

duburcqa commented Apr 3, 2026

Uh oh!

github-actions Bot commented Apr 3, 2026

Uh oh!

Lidang-Jiang commented Apr 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants