feat(training): RL parity — evaluate(), staged_reward, vectorized PPO, gym adapter (+ render output_path)

[cagataycali]

Kept as Draft on purpose. This is a decomposition tracker, not a mergeable unit. It still carries the full 2112-line mega-diff whose pieces are landing as focused children (#918, #920, queued follow-ups). Merging this branch as-is would undo the decomposition and conflict with its own children. Draft status makes that mechanically impossible until every piece has landed, at which point this PR is closed (not merged). Reopen-for-review only if the remaining follow-up content is ever reduced into a standalone reviewable diff.

Status: decomposing into focused PRs (not force-rebasing)

Per review disposition, this multi-concern branch is being landed as small, separately-reviewable units instead of a single mega-diff. The sim-bug fixes it originally carried (_target_quat vertical-camera +Y fallback, eval_policy sole-robot resolve, actuator-keyed policies) already landed independently on main, so they are dropped. The remaining net-new RL-parity work is being split:

Piece	Focused PR	Status
VecSimEnv (N-env threaded CPU/MuJoCo vec-env)	#918	open, CI green, awaiting review
BaseRLAlgo.evaluate() + load_checkpoint() + tests	#920	open, APPROVED
staged_reward declarative phase-machine + tests	follow-up	queued
GymSimEnv adapter + vectorized PPO rollout path	follow-up, maps to #863	queued
render(output_path=...) ergonomic	follow-up	queued

Follow-ups are opened one at a time to keep each diff surgical and individually validatable (and to respect parallel-PR limits). This PR stays open (as a Draft) as the decomposition tracker until its pieces have all landed; it will not be force-rebased into a single un-reviewable unit, and it will be closed rather than merged.

---

Original Summary (for reference)

The RL-training parity work: deterministic evaluate(), declarative staged_reward, SimEnv.n_substeps, a vectorized PPO rollout path, and the VecSimEnv / GymSimEnv adapters, plus an optional render(output_path=...) with a path-traversal guard. All net-new on top of the from-scratch RL trainers; none overlapping anything already on main.

§13 Review Round Changelog

Round	Concern	Fix Commit	Pin Test
R1	CI blocked: ruff + mypy errors introduced by new RL modules	6770922	N/A (lint-only fix; all 47 PR-specific tests pass)
R2	Branch is a multi-concern mega-diff that exceeded a sane review budget; camera fixes already on main. Decompose into focused PRs (#918, #920, queued follow-ups) rather than force-rebase.	no code push - decomposition tracker	per-piece in each focused PR
R3	Tracker sat in APPROVED+mergeable state - an accidental-merge footgun that would undo the decomposition. Converted to Draft so it is mechanically unmergeable; it will be closed (not merged) once children land.	no code push - state change only	N/A

[yinsong1986]

Summary

This PR lands two related bodies of work: four MuJoCo-sim correctness fixes and four RL-training parity additions. The sim fixes add robot_action_keys() (an ABC default mirroring joint names plus a MuJoCo override that returns namespace-scoped actuators, so policies on tendon-gripper / passive-joint robots stop silently no-op'ing), a world-+Y fallback in _target_quat for vertical cameras, an optional render(output_path=...) with a bad-char + ..-component guard on the LLM-supplied path, eval_policy auto-resolving the single robot to match run_policy, and a per-episode stateful-reward reset in DeclarativeBenchmark.on_episode_start. The RL parity work adds a declarative staged_reward phase machine (compiled through the existing closed predicate registry, no eval), BaseRLAlgo.evaluate() / load_checkpoint() / _deterministic_action(), SimEnv.n_substeps (default 5), a vectorized PPO rollout path gated on num_envs > 1, and the VecSimEnv / GymSimEnv adapters.

I walked the changed paths against the repo conventions. The render output_path guard rejects shell metacharacters and .. components on the user-supplied path; the vectorized GAE correctly broadcasts the scalar advantage recurrence over the (T, N) env axis and bootstraps truncations from the captured pre-reset terminal_obs; latest_checkpoint referenced from the base evaluate() is defined on both concrete trainers (ppo, fast_sac) so it resolves at runtime; the single executor in VecSimEnv is created once and reused per the AGENTS.md perf rule. The n_substeps=5 default is a training-dynamics tuning change rather than a wire/schema/CLI one-way door. New code carries no emojis or host paths and __all__ exports only public names. No blocking concerns found.

Verification suggestions

• hatch run test tests/training/test_rl_ppo_vectorized.py tests/training/test_rl_vec_env.py tests/training/test_rl_evaluate.py tests/training/test_rl_gym_env.py tests/simulation/test_staged_reward.py to confirm the RL parity suite is green in isolation.
• hatch run test tests/simulation/mujoco/test_robot_action_keys.py tests/simulation/mujoco/test_camera_vertical_lookat.py to confirm the embodiment action-key and vertical-camera regressions pass with MuJoCo present.

[cagataycali]

Maintainer disposition (empirically re-verified against main @ c991151)

This PR is APPROVED but has drifted far from main. I partitioned every changed file against current main and ran the render artifact in a CI-equivalent environment. Recommendation: don't force-rebase this 2266-line branch; instead land it as small, separately-reviewable pieces. Nothing here is being closed — this is guidance so you can land your work cleanly under your own approval.

1. Superseded — already on main (drop on re-land)

Four of the five functions in tests/simulation/mujoco/test_camera_vertical_lookat.py are redundant. main's tests/simulation/mujoco/test_spec_builder.py::TestTargetQuat already covers the _target_quat world-+Y fallback for top-down / straight-up / horizontal / degenerate cases — and goes further, compiling an MjSpec and asserting the camera's -Z forward axis actually points at the target (test_vertical_camera_forward_points_at_target). The feature itself (_target_quat +Y fallback, eval_policy sole-robot, send_action(n_substeps=), render output_path) is also already on main.

2. The ONE net-new, validated unit worth cherry-picking now

test_camera_vertical_lookat.py::test_topdown_camera_renders_object — the end-to-end render proof (build scene -> top-down camera -> count red pixels). main proves the quaternion geometrically but has no render-level proof; this closes that integration gap.

Verified this cycle on main: deterministic 1444 red px on 3/3 runs (~2.5s each).
Validation envelope (important): passes only under MUJOCO_GL=osmesa; it errors under egl/glfw/default. CI already sets MUJOCO_GL: osmesa (.github/workflows/test-lint.yml:90, with libosmesa6-dev), so it is CI-safe — but the file pins no backend, unlike main's render tests which all do os.environ.setdefault("MUJOCO_GL", "glfw"). Please align it to the existing convention (or guard with pytest.importorskip + an osmesa skip) before landing, so it doesn't become an env-coupled flake for local contributors.

Suggested minimal land: cherry-pick just that one function into a new tests/simulation/mujoco/test_render_topdown_camera.py (matches the test_render_* naming convention) — a ~30-line, zero-production-risk coverage gain.

3. Genuinely net-new FEATURE code — split into focused PRs

Absent from main, real value, but multi-concern and needs sim-stepping validation that shouldn't ride on a test-only diff:

• gym_env.py, vec_env.py, vectorized PPO -> map to roadmap feat(simulation): GymVecEnv adapter — (num_envs, *) step/reset/get_observations over MuJoCo + Newton + MJWarp #863, its own PR
• evaluate() + tests/training/test_rl_evaluate.py -> separate PR
• staged_reward + tests/simulation/test_staged_reward.py -> separate PR
• env.py n_substeps, gym_env/vec_env plumbing -> with the feat(simulation): GymVecEnv adapter — (num_envs, *) step/reset/get_observations over MuJoCo + Newton + MJWarp #863 PR

Why not just rebase this branch?

Force-rebasing reintroduces the full mega-diff, invalidates the existing approval, and bundles >=3 unrelated concerns into one un-reviewable unit. Splitting per (2)/(3) keeps each piece surgical and individually validatable.

Disposition only — no branches touched, nothing closed. The close/rebase/split call is yours.

[yinsong1986]

Summary

This PR adds RL-training parity on top of the from-scratch RL trainers (#877): BaseRLAlgo.evaluate()/load_checkpoint()/_deterministic_action() (deterministic eval peer of train(), frozen normalizer, side-effect-free train/eval mode restore, weights_only=True checkpoint load), a declarative staged_reward phase-machine predicate compiled through the closed registry (no eval), a vectorized PPO rollout path (num_envs>1 via VecSimEnv with per-env GAE bootstrapping from captured terminal_obs), a GymSimEnv gymnasium adapter with the correct terminated-vs-truncated split, SimEnv.n_substeps, and an optional render(output_path=...) with a path-traversal/shell-metachar guard. The diff matches the description: the sim-bug fixes that landed independently are dropped, and the remaining work is net-new RL parity.

I verified the changed paths against the codebase: n_substeps is part of the SimEngine.send_action ABC contract and the MuJoCo override, latest_checkpoint exists on both PpoTrainer/FastSacTrainer, the (T,N) GAE recurrence in compute_gae broadcasts correctly over the env axis, VecSimEnv reuses a single ThreadPoolExecutor (AGENTS.md perf rule), and the staged_reward factory stays inside the closed-registry/no-eval contract. The render(output_path) write is argv-style (open(...,'wb'), no shell) with no trusted-prefix join, so the ../shell-metachar guard is reasonable defense-in-depth rather than a traversal boundary. No emojis in tool/result/log strings, no host paths, ValueError guards on bad inputs. Lint is clean and the full RL/staged-reward unit suite passes locally (48 tests). The two prior CodeQL findings (unused register_predicate import, unused sb3 variable) are already resolved at head.

What's good

• Behaviour-preserving dispatch: num_envs==1 keeps the byte-identical single-env path; vectorization is opt-in.
• evaluate() restores pre-eval train/eval mode and freezes the normalizer, closing the BatchNorm/Dropout stat-freeze footgun.
• load_checkpoint uses weights_only=True, closing the legacy-unpickler RCE surface.
• Correct GAE bootstrapping: truncations bootstrap from the captured terminal obs, real terminals zero the bootstrap.
• Stateful-term reset() wired in both SimEnv.reset and DeclarativeBenchmark.on_episode_start so phase state never leaks across episodes.

Verification suggestions

• hatch run test -- tests/training/test_rl_vec_env.py tests/training/test_rl_ppo_vectorized.py tests/training/test_rl_evaluate.py tests/training/test_rl_gym_env.py tests/simulation/test_staged_reward.py
• GPU/MuJoCo host: hatch run test-integ -- tests_integ/training/test_so101_pick_place_e2e.py to confirm the reach-convergence assertion (>50% gripper->cube distance closure).

[yinsong1986]

Summary

This PR adds RL-training parity on top of the from-scratch RL trainers: a deterministic BaseRLAlgo.evaluate() (frozen normalizer via update=False, gradients off, side-effect-free train/eval-mode restore, weights_only=True checkpoint load) plus load_checkpoint()/_deterministic_action(); a declarative staged_reward phase-machine predicate compiled entirely through make_predicate (closed registry, no eval); SimEnv.n_substeps; a vectorized PPO rollout path that wraps N independent SimEnv in a reused-thread-pool VecSimEnv and bootstraps truncations from captured terminal_obs; a GymSimEnv gymnasium adapter that correctly splits terminated vs truncated; and an optional render(output_path=...) with a path-traversal/shell-metacharacter guard. The diff is additive: the num_envs != 1 PPO validation is loosened (single-env path is byte-for-byte unchanged), new symbols are net-new exports, and send_action(..., n_substeps=...) matches the existing SimEngine ABC signature so no call site breaks.

What's good

• evaluate() is correctly side-effect-free: it captures and restores the pre-eval .training flag for both the actor-critic and the normalizer, and freezes the normalizer with update=False so eval never shifts running stats.
• Stateful reward reset() is hooked in both SimEnv.reset and DeclarativeBenchmark.on_episode_start via duck-typed getattr(term, "reset", None), so per-episode phase state can't leak and stateless terms are untouched.
• VecSimEnv reuses a single ThreadPoolExecutor for its lifetime (AGENTS.md: never create executors in the hot loop) and fails loudly on heterogeneous sub-env dims; the autoreset path captures the true terminal obs before reset, which is load-bearing for GAE truncation bootstrapping.
• load_checkpoint uses weights_only=True, closing the legacy-unpickler RCE surface.
• Good scope discipline and dense regression tests pinning the eval schema, determinism, success-rate semantics, phase monotonicity, one-time bonus, and validation/safety rejections.

Verification suggestions

• hatch run test tests/training/test_rl_evaluate.py tests/training/test_rl_vec_env.py tests/training/test_rl_ppo_vectorized.py tests/simulation/test_staged_reward.py to exercise the CPU-only fakes.
• Spot-check the vectorized GAE path matches single-env on num_envs=1 (_collect_rollout_single vs _collect_rollout_vectorized should agree on identical seeds).

[yinsong1986]

Summary

This PR carries the RL-training parity work atop the from-scratch RL trainers: a declarative staged_reward phase-machine predicate compiled through the closed make_predicate registry (no eval), SimEnv.n_substeps plumbing plus per-episode reset() of stateful reward terms, a GymSimEnv gymnasium adapter that splits terminated vs truncated correctly at the source, a vectorized VecSimEnv-backed PPO rollout path (num_envs > 1), and an ergonomic render(output_path=...) with a ../shell-metacharacter path guard. As the description makes explicit, this branch is a Draft decomposition tracker, not a mergeable unit: its pieces are landing as focused children (#918 VecSimEnv, #920 evaluate(), queued follow-ups) and the tracker will be closed, not merged once they all land. Draft status makes accidental merge mechanically impossible.

Reviewing the diff itself: the staged_reward factory validates stages eagerly (non-empty list, allowed keys, predicate-call dicts, non-final-stage gate, numeric bonus) and stays inside the closed-registry safety contract; the monotonic phase machine and one-time-bonus semantics are pinned by dedicated CPU tests. Stateful-term reset() is wired into both SimEnv.reset and DeclarativeBenchmark.on_episode_start via duck-typed getattr, so stateless function terms are unaffected. The vectorized PPO path captures the pre-reset terminal_obs for correct truncation value-bootstrapping and flattens (T, N, ...) to (T*N, ...) for the shared minibatch update; evaluate() defensively routes through a single sub-env on a VecSimEnv. The render(output_path=...) guard rejects .. components and shell metacharacters before writing, consistent with the project's existing path-safety conventions. The substantive content was already reviewed and APPROVED on a prior commit; no new commit introduces a new instance of a previously-flagged bug class. No blocking concerns found.

What's good

• Scope discipline: the multi-concern mega-diff is being decomposed into individually-reviewable children rather than force-rebased, and the tracker is held as a Draft so it cannot be merged out from under the decomposition.
• staged_reward is authored as data through the closed predicate registry (no eval), matching the project's DSL-safety stance.
• Strong test coverage for the net-new pieces: phase advance / one-time bonus / monotonicity / reset, the gymnasium check_env conformance and terminated-vs-truncated split, and vectorized batch-shape / throughput-multiplier pins.
• terminated vs truncated is split at the adapter source (the SB3 TimeLimit footgun), so downstream learners get value-bootstrapping right for free.

Verification suggestions

• hatch run test tests/simulation/test_staged_reward.py tests/training/test_rl_ppo_vectorized.py tests/training/test_rl_gym_env.py to confirm the 47 PR-specific tests pass on the head commit.
• With GPU + MuJoCo weights: pytest tests_integ/training/test_so101_pick_place_e2e.py -v to exercise the full RL stack (staged_reward + SimEnv + evaluate + GymSimEnv/SB3 + vectorized PPO) against real SO-101 physics.