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

tests/test_energy.py

- 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

genesis/engine/entities/rigid_entity/rigid_entity.py

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

tests/test_rigid_physics.py

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

genesis/engine/entities/rigid_entity/rigid_entity.py

genesis/engine/solvers/rigid/rigid_solver.py

tests/test_rigid_physics.py

duburcqa · 2026-03-29T13:24:33Z

tests/test_rigid_physics.py

+@pytest.mark.required
+def test_energy_analytical_and_conservation(show_viewer, tol):
+    g = 9.81
+    dt = 0.01
+    h0 = 1.0
+    radius = 0.1
+    n_steps = 100
+    undamped_sol_params = [0.0, 0.0, 0.9, 0.95, 0.001, 0.5, 2.0]
+
+    scene = gs.Scene(
+        sim_options=gs.options.SimOptions(dt=dt, gravity=(0, 0, -g)),


Parameterize this test with both Euler and approximate_implicitfast integrators.

Done in 47129f5. Parameterized with both Euler and approximate_implicitfast. Also changed timeconst from 0 to 0.02 in undamped sol_params — the previous timeconst=0, dampratio=0 combination caused NaN constraint forces on the field backend CI.

Add a FIXME comment in the code to document this issue.

duburcqa · 2026-03-29T13:25:10Z

tests/test_rigid_physics.py

+        ke_a.append(sphere_a.get_kinetic_energy().item())
+        pe_a.append(sphere_a.get_potential_energy().item())
+        ke_b.append(sphere_b.get_kinetic_energy().item())
+        pe_b.append(sphere_b.get_potential_energy().item())
+        z_a.append(sphere_a.get_pos()[..., 2].item())
+        z_b.append(sphere_b.get_pos()[..., 2].item())


Remove .item(), it is useless.

duburcqa · 2026-03-29T13:29:23Z

tests/test_rigid_physics.py

+    # First impact timestep (sphere center reaches ~radius above ground)
+    impact_step = next((i for i, z in enumerate(z_a) if z <= radius + dt), n_steps)


Use real collision detection for this, ie

impact_step = -1 for i in range(n_steps): [...] if impact_step < 0 and scene.rigid_solver.collider._collider_state.n_contacts.to_numpy().any(): impact_step = i assert impact_step > 0

It is relying on private API but it is fine in this context.

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 requested changes Mar 28, 2026

View reviewed changes

duburcqa reviewed Mar 28, 2026

View reviewed changes

tests/test_energy.py Outdated Show resolved Hide resolved

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

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 28, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 29, 2026

View reviewed changes

tests/test_rigid_physics.py Show resolved Hide resolved

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

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 29, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 29, 2026

View reviewed changes

genesis/engine/entities/rigid_entity/rigid_entity.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 29, 2026

View reviewed changes

genesis/engine/solvers/rigid/rigid_solver.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 29, 2026

View reviewed changes

tests/test_rigid_physics.py Outdated Show resolved Hide resolved

duburcqa reviewed Mar 29, 2026

View reviewed changes

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

		# First impact timestep (sphere center reaches ~radius above ground)
		impact_step = next((i for i, z in enumerate(z_a) if z <= radius + dt), n_steps)

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!

duburcqa Mar 29, 2026

Choose a reason for hiding this comment

Uh oh!

Lidang-Jiang Mar 29, 2026

Choose a reason for hiding this comment

Uh oh!

duburcqa Mar 29, 2026

Choose a reason for hiding this comment

Uh oh!

duburcqa Mar 29, 2026

Choose a reason for hiding this comment

Uh oh!

duburcqa Mar 29, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants