[MISC] Add lazy velocity update to kinematic entity.

[jeetrex17]

Description

Defers link-velocity propagation for kinematic entities until it is needed.

set_dofs_position fuses the degree-of-freedom write, qpos synchronization, position forward kinematics, and center-of-mass propagation into one kinematic kernel. get_links_vel, get_links_ang, and set_dofs_velocity refresh velocities when needed. A velocity write following a skipped pose update refreshes pose and velocity together so the resulting link velocities remain consistent. Other pose setters retain their eager full refresh.

Measured independent host-side kernel launches per steady-state step:

Workflow	main	this PR
write positions, read positions	2	1
write positions, read positions + velocities	2	2

The fused kernel contains internally sequenced top-level loops; the reduction is in independent host-side launches.

Two consistency fixes are included:

• get_state refreshes link poses before capture, keeping them consistent with the generalized coordinates stored alongside them.
• RigidSolver.get_links_ang remains a direct read because rigid stepping owns velocity propagation.

Related Issue

Supersedes #2558.

Motivation and Context

Kinematic entities are commonly driven by pose setters while many consumers never read link velocities. The position-only path therefore paid for unnecessary propagation. Fusing the position write with position forward kinematics reduces that path from two independent launches to one without increasing the position-plus-velocity path.

Benchmarks were run on Ubuntu 24.04 with an RTX 2080 Ti at 20k environments, using a 45-second warmup and 15-second recording window across five interleaved A/B pairs.

Scenario	main	this PR	Change
Position only	32.56M env updates/s	38.08M env updates/s	+17.0% (95% CI +12.2% to +21.8%)
Position followed by link-velocity read	29.81M env updates/s	31.57M env updates/s	+5.9% (95% CI +4.4% to +7.4%)

Every interleaved pair was positive in both scenarios. These results were collected with temporary focused fixtures for the two access patterns above. The fixtures were removed from the final diff after review because they were only needed to analyze this PR. The previous mixed fixture measured +0.67% because unrelated rigid-scene work diluted the optimized kinematic path.

Getter-triggered velocity refreshes run over the full batch. Only envs_idx=None is treated as a definitely complete batch for freshness tracking because inspecting an arbitrary device mask would require synchronization.

How Has This Been / Can This Be Tested?

pytest tests/rigid/test_kinematics.py::test_setters tests/rigid/test_kinematics.py::test_link_velocity -v
pytest tests/rigid/test_kinematics.py -q

The full kinematics file passes locally: 19 passed and 2 pre-existing macOS skips on an Apple M-series CPU with Python 3.13. The stale-pose velocity regression fails on the previous head with a maximum absolute link-velocity difference of 0.62 and passes after the fix. Deterministic instrumentation confirms one launch for position-only workflows and two for position-plus-velocity workflows, including the boolean-mask setter path.

Checklist:

• I read the CONTRIBUTING document.
• I followed the Submitting Code Changes section of CONTRIBUTING document.
• I tagged the title correctly (including BUG FIX/FEATURE/MISC/BREAKING)
• I updated the documentation accordingly or no change is needed.
• I tested my changes and added instructions on how to test it for reviewers.
• I have added tests to cover my changes.
• All new and existing tests passed.

[duburcqa]

The key is to reduce the number of independent kernel calls, or at the very least, make sure that it does not increase. Otherwise, it is likely to cause performance regression instead of improvement.

By the way, since this is a perf improvement PR, it would be nice to attach some benchmark stats to show the benefits of lazy velocity update when only position is requested (to show improvement), and when position+velocity is requested (to show no regression).

[jeetrex17]

Thanks. I reworked this around the number of independent host-side kernel launches.

The measured steady-state counts are now:

• Position only: 2 to 1
• Position plus velocity: 2 to 2

set_dofs_position now fuses the DOF write, qpos synchronization, position FK, and COM propagation into one kinematic kernel. The kernel contains internally sequenced loops, but only one independent host-side launch.

I also ran interleaved local Metal A/B benchmarks with 20k environments, but the machine could not be isolated from background load. The resulting 16% to 35% run-to-run ranges make the figures inconclusive, so I am not quoting a throughput number. The +12.43% figure in the earlier description was measured on the previous design and has been removed.

I restored the anymal_uniform_kinematic benchmark removed in #3249 so the position-only workflow can be measured on reference hardware.

[duburcqa]

I also ran interleaved local Metal A/B benchmarks with 20k environments, but the machine could not be isolated from background load. The resulting 16% to 35% run-to-run ranges make the figures inconclusive, so I am not quoting a throughput number.

@jeetrex17 We will need reliable numbers before merging.

[duburcqa]

We are not going to merge this, but it is fine for analysing the benefits of this PR temporarily.

[jeetrex17]

@jeetrex17 We will need reliable numbers before merging.

I do not have regular access to CUDA hardware, so after the local Metal results were inconclusive, I was able to arrange temporary access to a workstation in my university lab and rerun the benchmarks on CUDA

I ran five interleaved A/B pairs with the official benchmark runner on Ubuntu 24.04 with an RTX 2080 Ti and 20,000 environments. Each sample used a 45-second warmup and 15-second recording window. Throughput is measured in kinematic environment updates per second.

Scenario	main	this PR	Change
Position only	32.56M	38.08M	+17.0% (95% CI: +12.2% to +21.8%)
Position followed by link-velocity read	29.81M	31.57M	+5.9% (95% CI: +4.4% to +7.4%)

Every interleaved pair was positive in both scenarios.

These results were collected with temporary focused fixtures for the two access patterns. Per your review, the fixtures have been removed from the final diff, so the benchmark suite is unchanged. The full methodology and results are in the PR description.

[duburcqa]

This seems completely unrelated and could have unexpected side effects.

[duburcqa]

I don't like this kind of tricks. It is just a workaround. It seems that 'is_forward(pos|vel)_updated' should be refactored to be boolean masks.

[duburcqa]

same.

[duburcqa]

same.

[duburcqa]

same.

[duburcqa]

All these arguments fit on 1 line.

[duburcqa]

Add a code comment to explain what this is checking. Moreover, we should only assert what is part of the specification of genesis. From what I see, it is not immediately clear whether this is spec or side-effect.

[duburcqa]

Clone is useless. All genesis getters are returning by value.

[duburcqa]

Clone is useless. All genesis getters are returning by value.

[duburcqa]

(False, True) is not supported? In any case, specifying 'dtype=torch.bool' is useless.

[duburcqa]

Keep as a native python nested tuple.

[duburcqa]

This formatting is quite ugly and hard to read.

[duburcqa]

Move test constant on top the of unit test.

[duburcqa]

just do scene.kinematic_solver

[duburcqa]

I'm skeptical so many checks are necessary to assert what you want.

[duburcqa]

We do not assert shape. it is trial. We are implemented integration tests, not unit tests in practice. They are here to catch real bugs. if something is always valid by design or the likelihood to cause any other test to fail indirectly is high, when the assertion is skipped, to reduce the CI cost and improve the signal-to-noise ratio.

[duburcqa]

What is this supposed to check?

[duburcqa]

Add 1-line comment to each block

[jeetrex17]

Thanks for the detailed review. I’ve addressed the latest comments, simplified the tests, added coverage for the masked paths, and clarified the skip_forward behavior. The targeted link-velocity test passes.

[duburcqa]

It turns out setting 'self._is_forward_vel_updated' in python-scope using torch is faster than passing torch tensors directly to a quadrants kernels... Sorry about that. It should probably be mentioned in CODING_GUIDELINES.md.

[jeetrex17]

I’ll move the freshness-mask updates currently performed inside Quadrants kernels to Python.

[duburcqa]

Now that I'm think about it, these masks should be used to skip update INSIDE the kernels for the envs that are already up-to-date. This means that this mask as to be passed one way or the other...

I'm afraid getting this PR right is harder than it seems if we want to address all the issues and edge-cases without compromising performance.

[duburcqa]

Why are both needed? I understand you point about GPY sync issue, but this is only if you need to do something like if self._is_forward_vel_updated_for_all_envs.all():. But if a mask is used everywhere, this should no longer be necessary, removing the need for 'all_envs' variant.

[jeetrex17]

The host summaries are only a fast path for fully fresh batches. Without them, getters would either need a synchronizing .all() check or launch a no-op refresh kernel even when every environment is already fresh, which would lose the 2 to 1 launch reduction on the position-only path. The per-environment masks are still needed for mixed batches, so the two serve different purposes.

[duburcqa]

Sure. But computing '_is_forward_vel_updated_for_all_envs' requires a GPU sync anyway? it is only free when the user does not specify any envs_idx mask. This means that this shortcut only every take effect when envs_idx is not specify. I don't think it is a good idea, as per-env operations are becoming more and more common, and per-env stepping is going to be available soon.

In what extend this does affect performance? If updating position-only, you can still skip velocity update by default, which was the first objective of this PR.

[jeetrex17]

You are right that the position-only setter keeps its one-launch improvement without the host summaries. My earlier statement about losing that reduction was inaccurate.

I measured the current PR against a local variant with both host summaries removed on an RTX 2080 Ti, using 20k environments, five interleaved A/B pairs, a 45-second warmup, and a 15-second recording window:

Workflow	Current PR	Masks only	Change
Position + scene.step()	32.47M	27.78M	-14.42% (95% CI -16.17% to -12.67%)
Repeated fresh velocity read	348.55M	151.20M	-56.62% (95% CI -57.53% to -55.71%)
Position only	33.87M	34.00M	+0.39% (95% CI -1.04% to +1.81%)
Position + velocity read	24.72M	24.63M	-0.38% (95% CI -1.39% to +0.63%)

The first two regressions occurred in all five pairs. The latter two differences are statistically inconclusive.

The host summaries are not computed from the device masks with .all(). They are plain Python booleans maintained conservatively from the call scope: a known full-batch refresh can set them to True, while any subset operation sets them to False. The per-environment masks remain the source of truth and are still passed into the kernels to skip fresh environments.

Without the summaries, substep_post_coupling and already-fresh getters launch a no-op refresh kernel. Under sustained per-environment operations, the summaries simply remain False and the code behaves like the mask-only variant until a full-batch refresh re-arms them.

Based on these results, I think we should keep both: the masks for per-environment correctness and skipping, and the host booleans only as a full-batch launch gate.