[Warp] Restore --frontend warp on Reach and gate the supported task set in CI - #6900
Conversation
The Reach task configs gained `rewards.success` and `terminations.success`, which resolve to `is_terminated_term` and `pose_command_success`. Neither had a warp counterpart, so `--frontend warp` rejected Isaac-Reach-Franka and Isaac-Reach-UR10 at env construction. `is_terminated_term` is a ManagerTermBase subclass rather than a plain function, so the twin is a class too: the warp RewardManager instantiates it at _prepare_terms and calls the instance. It resolves the term selection to a column mask once at init and reads the per-term done buffer, which is exposed as a new `term_dones_wp` property. `pose_command_success` recomputes the pose error in a kernel instead of copying the stable term's result. `UniformPoseCommand._compute_success` allocates a fresh tensor per call, so its pointer moves between steps; a captured CUDA graph records source pointers at capture time and would replay against a stale buffer. The command's sticky success tracker stays correct because `_update_metrics` already ORs the same value in on every step. Parity tests cover both twins against the stable implementations, eager and under graph capture. The pose-command fixture derives its command from each body's actual pose plus a known error so the errors straddle the thresholds -- a command sampled independently misses every threshold and compares all-False, which passes against almost any kernel.
The list of tasks that `--frontend warp` accepts was maintained by hand, so it only covered a task once someone remembered to add it. Sweeping the gym registry instead shows the record had already drifted: four contrib velocity tasks adapt cleanly and were never listed. The sweep runs `check_compatibility` over every registered stable manager-based task. That is pure config-tree work with no simulator, so all 104 resolve in about ten seconds. The checked-in set is now a regression pin rather than an allowlist. Losing a task always fails. Gaining one fails too, but only when every candidate config was importable, since an environment missing optional dependencies cannot see the full set; those tasks are named in the skip message rather than silently dropped. `check_compatibility` calls `adapt_cfg` and returns the reason instead of raising, so the verdict cannot drift from what the frontend actually does.
Every test job either names a package with a positive filter or, for the core jobs, excludes anything matching `isaaclab_`. No job named this package, so all nine of its test files fell between the two forms and were collected by nothing. Replaying the filter logic over the tree shows test files matched by no job dropping from thirteen to four; the remainder belong to isaaclab_ppisp and are left for a separate change.
Resolving the command term, its thresholds, the tracked body index and the warp view of the command buffer per call left non-capturable work on the hot path, including a hasattr-guarded warm-up cache on the function object. All of it is fixed for the run, so it belongs in __init__; __call__ is now a single kernel launch. The command manager is constructed before the termination manager, so the lookup is available at init.
A two-line ternary used twice does not need a named helper.
Three divergences from the stable term, each with a regression test that fails
when the change is reverted.
The sticky per-episode success bit was never set. The stable term ORs it inside
compute_success(), which runs before the episode is reset; the twin only wrote
the termination output, and _update_metrics() runs after reset() has already
read and cleared the tracker, so a terminating step went unrecorded. The write
now happens in the kernel rather than after the launch, because the termination
manager runs graph-captured by default and host-side work between launches is
not replayed.
The threshold comparisons were reject predicates ("veto when error >= bound"),
which leave a non-finite error reported as success -- terminating and rewarding
a diverged environment. They are accept predicates now, matching the stable
`error < bound` and failing closed.
A negative threshold was encoded as "not configured". UniformPoseCommandCfg
accepts any float without validation, so a configured negative bound was
silently ignored instead of denying success. Configuredness is carried as its
own flag.
Thresholds stay resolved at init; the docstring now records that a curriculum
mutating them at runtime is not honoured, since a kernel scalar is baked into
the graph at capture.
Managers run graph-captured by default, so a term that is not capture-safe is on the normal execution path, not an edge case. Capture-and-replay alone does not prove safety: replaying against unchanged memory re-reads the same valid bytes, so a term holding a stale pointer still looks correct. Each term here is captured, has its inputs overwritten in place, then replayed and compared against the stable implementation on the new data. Terms are discovered from the warp MDP modules rather than listed, so a new term without a declaration fails the gate. Declaring one takes a CaptureSpec naming its modality and a builder that wires up the environment, parameters and in-place mutation; the runner is shared. Enumerating also separated three groups that convention had blurred: terms already covered by the hand-written mutation classes, three observation terms with no capture coverage anywhere, and eight event terms that cannot be captured at all because they launch over a variable env_ids subset at reset cadence.
Isaac-Ant and Isaac-Humanoid stopped adapting under --frontend warp when the locomotion reward cfg gained `terminated_penalty` and `survival_success_rate`; neither had a warp counterpart, which the registry sweep reports as lost support. `survival_success_rate` follows the cartpole twin: counts are accumulated on device and the rate is exposed as a persistent tensor view through the reward manager's reset extras, so no host readback happens where the stable term calls .item(). The capture-safety gate now also walks the per-task warp mirrors under isaaclab_tasks_experimental, not just the shared modules, so a twin added to a task mirror is held to the same declaration rule. Enumerating them shows that package has no test directory at all: all twenty pre-existing per-task twins are recorded as having no coverage rather than assumed fine. CaptureSpec gains expect_nonzero so a term whose output is constant by design -- survival_success_rate writes zeros, its real work being in reset() -- can be exercised without the non-degeneracy guard rejecting it.
2*acos(|w|) reads the scalar part as a cosine, which holds only for a unit quaternion and is ill-conditioned as the error approaches zero: w = cos(theta/2) carries theta only at second order, so in float32 the value rounds to exactly 1 and the reported angle collapses to 0 below roughly 1e-4 rad. Measured against the analytic angle, 2*acos is 2.3% low at 1e-3 rad and returns zero at 1e-4. 2*atan2(|xyz|, |w|) recovers the angle from the ratio, where the numerator carries theta at first order. It is exact across the same range and independent of the quaternion's norm. This matters most for orientation_command_error, a continuous reward whose error shrinks toward zero as a policy converges -- precisely where acos stops resolving. pose_command_success gates on a threshold three orders of magnitude above the affected range, but computes the same quantity, so both use one form. survival_success_rate no longer substitutes 0.0 when a reset selects no environments: the stable term's mean over an empty selection is NaN, and publishing a real 0% sample misreports it as a measurement.
Both gates could report success while the property they guard was violated. Terms were keyed by bare name, but survival_success_rate is defined by both the cartpole and locomotion mirrors, so a spec for one marked the other declared. Identity is now module-qualified throughout discovery, specs and the backlog. Discovery skipped action terms even though ActionManager is graph-captured like every other manager, leaving three of them unable to fail the declaration check. The event-term rows claimed a variable env_ids launch dim made them uncapturable. That is not what they do -- every one launches over env.num_envs with a boolean mask, which is the capture-safe form -- and four already have capture coverage in the events parity file. The rows now say which are covered and which are simply untested. The runner captured only each term's output call, so a stateful write moved to host code would replay as a no-op unnoticed. Terms may now assert their side effects after replay, and that assertion runs before the stable reference is called: the stable terms write the same shared state they are compared against, so calling the reference first repaired exactly what the assertion checks. The registry sweep matched only string cfg entry points and silently omitted any other supported form; it now loads through load_cfg_from_registry.
Greptile SummaryThis PR restores Warp frontend support for Reach, Ant, and Humanoid tasks while adding registry-derived compatibility and CUDA-graph capture-safety coverage.
Confidence Score: 4/5The PR appears safe to merge functionally, with one non-blocking CI supply-chain hardening issue to address. The Warp twins and their coverage align with the stable manager contracts examined; the remaining concern is the mutable checkout action reference in the newly added self-hosted CI job. Files Needing Attention: .github/workflows/build.yaml
|
| Filename | Overview |
|---|---|
| source/isaaclab_experimental/isaaclab_experimental/envs/mdp/terminations.py | Adds a capture-safe pose-command success twin with stable-compatible pose checks and sticky success tracking. |
| source/isaaclab_experimental/isaaclab_experimental/envs/mdp/rewards.py | Adds a selected-termination reward twin backed by the manager's ordered per-term done buffer. |
| source/isaaclab_tasks_experimental/isaaclab_tasks_experimental/core/locomotion/mdp/rewards.py | Adds graph-capturable termination-penalty and survival-success metric twins. |
| source/isaaclab_experimental/test/envs/mdp/test_capture_safety.py | Adds discovery-based capture declarations and mutated-input replay parity checks for Warp MDP terms. |
| source/isaaclab_experimental/test/envs/test_frontend_cfg_conversion.py | Replaces manually selected adaptation coverage with a registry sweep and pinned supported-task regression set. |
| .github/workflows/build.yaml | Adds CI coverage for isaaclab_experimental, but the new self-hosted job uses a mutable checkout action reference. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
Registry[Registered manager-based tasks] --> Sweep[Warp compatibility sweep]
Sweep --> Adapt[WarpFrontend.adapt_cfg]
Adapt --> Twins[Warp MDP twins]
Twins --> Managers[Graph-captured managers]
Managers --> Replay[Capture and mutated-input replay tests]
Replay --> Stable[Stable-term parity comparison]
Reviews (1): Last reviewed commit: "Merge remote-tracking branch 'upstream/d..." | Re-trigger Greptile
| github.event_name != 'push' && | ||
| needs.build.result == 'success' | ||
| steps: | ||
| - uses: actions/checkout@v6 |
There was a problem hiding this comment.
Mutable checkout action reference
The new self-hosted CI job executes actions/checkout@v6, whose tag can be retargeted without a repository change, exposing the runner to unreviewed action code. Pin this dependency to a full commit SHA so the executed workflow remains immutable. How this was verified: The added job directly invokes actions/checkout using the mutable v6 tag on a self-hosted runner.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
Isaac Lab Review Bot
The Warp frontend restoration is internally consistent: the new MDP twins, registry-derived support sweep, capture/replay coverage, CI package job, exports, and changelog fragments align with the affected execution paths. No candidate finding establishes a current contract violation requiring a pre-merge fix.
- Design and architecture: The frontend model remains coherent: check_compatibility delegates to the actual in-place adaptation path, and the new class terms follow the existing preallocated-output and graph-captured manager patterns. The pose-success twin intentionally aliases the stable command term's persistent private success buffer; this is a maintenance coupling, but the implementation documents the allocation assumption and tests the sticky write through replay, so the candidate does not demonstrate a present failure.
- API: WarpFrontend.check_compatibility and TerminationManager.term_dones_wp are additive and documented, including mutation and buffer-order semantics. The locomotion symbols are exported through the package stub. The is_terminated_term call signature preserves the manager-term parameter shape while resolving its selection at initialization, consistent with the class-term configuration model; the candidate did not establish a broken consumer or changed stable contract.
- Implementation: The changed producer and consumer paths were traced for per-term termination columns, timeout masking, pose-error calculation, sticky success tracking, locomotion metrics, task discovery, and CUDA-graph replay after in-place input mutation. The initialization-time freezing of pose thresholds is a disclosed behavioral tradeoff under capture rather than an unreported defect.
No blocking issues. No inline issue met the actionable-evidence threshold; the assessment above records the review feedback.
Automated review; human maintainers own approval decisions.
| kernel=_is_terminated_term_kernel, | ||
| dim=env.num_envs, | ||
| inputs=[ | ||
| env.termination_manager.term_dones_wp, |
There was a problem hiding this comment.
Let's avoid the _wp suffiex (it might be already there from before but it's cleaner without it)
|
|
||
|
|
||
| @wp.kernel | ||
| def _is_terminated_term_kernel( |
There was a problem hiding this comment.
This could actually be a 2D kernel, or even a tiled kernel. The tile is of shape tile_mask, and we can use tile_sum() or a tile_map().
(You can ignore this comment for now, I noticed that many of the kernels could be written in a much more optimized way, which we can tackle later)
| des_b = wp.vec3f(cmd[i, 0], cmd[i, 1], cmd[i, 2]) | ||
| des_w = root_pos_w[i] + wp.quat_rotate(root_quat_w[i], des_b) | ||
| cur_w = body_pos_w[i, body_idx] | ||
| dx = cur_w[0] - des_w[0] |
| # carry "configured" as its own flag: a threshold is any float, so no value can encode absence | ||
| position_threshold = command.cfg.position_success_threshold | ||
| orientation_threshold = command.cfg.orientation_success_threshold | ||
| self._check_position = position_threshold is not None |
There was a problem hiding this comment.
One thing you can do here and elsewhere is that thing like check_position can actually be used as wp.static inside the kernel definition. To do that you would need to define the kernel here inside init or in call.
|
|
||
| rng = np.random.RandomState(seed) | ||
| data = articulation.data | ||
| root_pos = data.root_pos_w.torch.cpu().numpy() |
There was a problem hiding this comment.
This looks weird. Whats wrong with root_pos_w.numpy()?
| return term | ||
|
|
||
|
|
||
| class MockPoseCommandManager: |
There was a problem hiding this comment.
HIgh level question, why do we need mocks? I prefer that tests use the original code if it is cheap to do so. The only exception are really expensive classes like Newton Model
| rng = np.random.RandomState(seed) | ||
| dones_np = rng.rand(num_envs, len(self.active_terms)) < 0.3 | ||
| self.term_dones = torch.tensor(dones_np, dtype=torch.bool, device=device) | ||
| self.term_dones_wp = wp.from_torch(self.term_dones) |
There was a problem hiding this comment.
Proxy Arrays here and elsewhere?
| # Stable task ids that adapt cleanly under ``--frontend warp``, as produced by | ||
| # :func:`_sweep_warp_support`. Do not curate this by hand: when the test fails it prints the | ||
| # exact set to paste back. | ||
| _WARP_SUPPORTED_TASKS = frozenset( |
There was a problem hiding this comment.
Any way we can track the warp supported envs automatically?
| return "{\n" + "".join(f" {task_id!r},\n" for task_id in sorted(task_ids)) + "}" | ||
|
|
||
|
|
||
| def test_warp_supported_task_set_matches_the_registry(): |
There was a problem hiding this comment.
Ignore above comment, it looks like we already are searching automatically. So why do we need to compare to a hardcoded list
There was a problem hiding this comment.
The goal is to enforce support. If it's resolved dynamically, it also means that torch mdp changes could break warp support. So the goal is to gradually add all parity mdps. And if we have all the envs supported, we can flip the test to ensure all envs are supported.
|
|
||
|
|
||
| @wp.kernel | ||
| def _survival_rate_kernel( |
There was a problem hiding this comment.
This is not even a real kernel?
Addresses review comments on #6900. - The pose-command parity fixture reached numpy via `.torch.cpu().numpy()`, routing warp buffers through a torch view and a host copy to get data warp hands back directly. Now uses `.warp.numpy()`. - `UniformPoseCommand` imports moved to the top of the file; they were local without a cycle to justify it. `ProxyArray.numpy()` is not equivalent and was tried first: unknown attributes forward to the torch view, so it inherits torch's CPU-only restriction and raises on the cuda tensors these fixtures use. Not included, raised on #6900 and better handled separately: - Dropping the `_wp` suffix. `term_dones_wp` belongs to a family of six (`time_outs_wp`, `dones_wp`, `terminated_wp`, `_truncated_wp`, `_terminated_wp`, `_scratch_term_mask_wp`); renaming one leaves the package less consistent than it is now. A rename of the family, with deprecations, is its own change. - `wp.static` for the resolved threshold flags, and 2D/tiled rewrites of the reward kernels. Both are optimizations that change how the kernels are defined and warrant their own benchmarking. ## Release backport - [x] <!-- backport-active-release --> Backport this pull request to the active release branch after it merges into `develop` Test plan: - [x] `pytest source/isaaclab_experimental/test/` — 205 passed, 1 skipped - [x] `uv run isaaclab -f` clean
Addresses review comments on #6900. - The pose-command parity fixture reached numpy via `.torch.cpu().numpy()`, routing warp buffers through a torch view and a host copy to get data warp hands back directly. Now uses `.warp.numpy()`. - `UniformPoseCommand` imports moved to the top of the file; they were local without a cycle to justify it. `ProxyArray.numpy()` is not equivalent and was tried first: unknown attributes forward to the torch view, so it inherits torch's CPU-only restriction and raises on the cuda tensors these fixtures use. Not included, raised on #6900 and better handled separately: - Dropping the `_wp` suffix. `term_dones_wp` belongs to a family of six (`time_outs_wp`, `dones_wp`, `terminated_wp`, `_truncated_wp`, `_terminated_wp`, `_scratch_term_mask_wp`); renaming one leaves the package less consistent than it is now. A rename of the family, with deprecations, is its own change. - `wp.static` for the resolved threshold flags, and 2D/tiled rewrites of the reward kernels. Both are optimizations that change how the kernels are defined and warrant their own benchmarking. ## Release backport - [x] <!-- backport-active-release --> Backport this pull request to the active release branch after it merges into `develop` Test plan: - [x] `pytest source/isaaclab_experimental/test/` — 205 passed, 1 skipped - [x] `uv run isaaclab -f` clean (cherry picked from commit 318abfd)
1. Summary
--frontend warpfails at env construction forIsaac-Reach-FrankaandIsaac-Reach-UR10:rewards.successandterminations.successresolve tois_terminated_termandpose_command_success, which had no warp twins. Adds both, matched to the stable terms and validated by training.source/isaaclab_experimental/had no CI job at all: all 9 of its test files were collected by nothing. Test files matched by no job: 13 → 4.Isaac-Ant/Isaac-Humanoid, which lost warp support on develop when the locomotion rewards gainedterminated_penaltyandsurvival_success_ratewithout twins — the sweep in this PR is what surfaced it.2. Why the existing gate did not catch it
test_frontend_cfg_conversion.pyalready asserted both Reach tasks adapt, and has since #5504 — it just never ran. Jobs either name a package with a positivefilter-pattern, or passnot isaaclab_, whichrun_tests.shconverts into an exclude onisaaclab_. This package matched neither form.3. Twin semantics
pose_command_successrecomputes the pose error in a kernel rather than copying the stable result, whose buffer is reallocated per call and would go stale under graph replay. Three points where a naive twin diverges from the stable term, each with a regression test:compute_success(), before resetreset()reads and clears first, soMetrics/success_ratereads 0error < boundis false → denies successerror >= boundis false → reports successis not None) → denies successThe sticky write happens inside the kernel, not after the launch: host-side work between launches is not replayed under capture.
Thresholds are resolved at init. A curriculum mutating them at runtime is not honoured, and cannot be while the manager is captured — a kernel scalar is baked into the graph. Recorded in the class docstring.
4. Capture-safety gate
Terms are discovered from the shared warp MDP modules and every per-task mirror under
isaaclab_tasks_experimental, rather than listed, so a new term without a declaration fails the gate. Enumerating separated three groups convention had blurred: terms already covered by the hand-written mutation classes, three observation terms with no capture coverage anywhere, and eight event terms that cannot be captured at all (variableenv_idslaunch dim at reset cadence).Capture-and-replay alone does not prove safety — replaying against unchanged memory re-reads the same valid bytes. Injecting a stale-view bug (
wp.cloneinstead of a zero-copy view) fails onlytest_term_is_capture_safe[pose_command_success]; all 17 other tests pass.5. Test plan
pytest source/isaaclab_experimental/test/— 201 passed, 1 skipped--frontend warp presets=newton_mjwarp:Metrics/success_rateIsaac-Reach-FrankaIsaac-Reach-UR10With the sticky write reverted, Franka trains identically (episode length still falls 360 → 252, so terminations fire) but
Metrics/success_ratestays 0.0000 for all 150 iterations — the metric bug in isolation.uv run isaaclab -fclean6. Out of scope
source/isaaclab_ppisp/has the same CI hole (4 uncollected test files).joint_pos_rel,joint_vel_rel,body_incoming_wrench.isaaclab_tasks_experimentalhas no test directory, so its twenty pre-existing per-task twins have no parity or capture coverage anywhere. They are recorded in the gate's backlog rather than assumed correct.2*acos(|w|)where the stable path normalizes first; the two differ by ~7e-4 rad at float32, far below any configured threshold. Pre-existing inorientation_command_error; left consistent rather than changed in one place.