Split out direct, SHAC, profiling, and accuracy extras - #61
Open
vmoens wants to merge 55 commits into
Open
Conversation
The bracket swap conditions were incorrect, causing the humanoid solver to diverge from MuJoCo C with default iterations=1. Three fixes: 1. Replace ad-hoc swap conditions with in_bracket (matching C's updateBracket / MJX), which only accepts candidates that narrow toward zero derivative. 2. Add one-sided phase emulation: when both bracket endpoints have same-sign derivatives (not yet bracketed), also accept candidates closer to zero, matching C's one-sided Newton phase. 3. Add early convergence check after initial Newton step, and add cross-bracket candidate swaps (6 candidates like MJX). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Raise _INLINE_CHOLESKY_MAX_SIZE from 16 to 32 so humanoid (nv=28) uses the inline path instead of torch.linalg.cholesky - Add clamp_min(s, 1e-12) to inline diagonal pivots to prevent NaN from numerically non-SPD mass matrices (degenerate physics states) - Add eps*I regularisation to the fallback path (nv > 32) as safety net - All changes are torch.compile friendly (no control flow/error handling) Fixes Cholesky crash in SAC humanoid training when envs reach extreme joint configurations that make the mass matrix non-positive-definite. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Replace data-dependent tensor shape checks (d.efc_J.numel() == 0) with static Python int checks from constraint_sizes(m) so torch._dynamo traces once without placing guards on tensor shapes. Also ensure make_constraint always uses the static max nefc for the output tensor, matching the MJX approach of fixed-size constraint arrays. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
The benchmark shows 200K+ steps/sec at B=2048 but SAC was getting 4K. Root cause: compile wrapped the inner step_fn, then vmap() was called on it every iteration. This prevented compile from fusing the batched kernels. Now we pre-vmap the function and compile the vmapped version, matching the benchmark's approach. Also caches the vmapped function instead of recreating it each call. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Frame skip loop (5 iterations) is now inside the compiled function, so torch.compile can fuse all 5 physics steps into one kernel launch - Use update_(ctrl=ctrl) instead of replace(ctrl=ctrl) to avoid shallow-copying the 80-field MjData dict every step - Cache vmap function instead of recreating per call Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Replace Python loop of n individual clone() + torch.stack with a single expand(n).clone() (one broadcast copy kernel). Use in-place add_ for noise. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Humanoid zoo env wrapping MuJoCo humanoid model - Direct backprop training script (humanoid-specific + generic) - Gradient accuracy tests for humanoid physics Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
When auto_reset=True, done envs are reset inside _step before building observations. This eliminates the separate maybe_reset → _reset path and allows the reset to be part of the same execution flow as physics. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
With _skip_maybe_reset=True, InitTracker never set is_init on episode boundaries, causing SACLoss value estimation to produce NaN. Instead, keep maybe_reset running but make _reset a no-op when auto_reset is enabled (since _step already handled the reset). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
The obs returned by _step must match the state that produced the reward. Previously, we reset done envs before building obs, so the replay buffer stored (terminal_reward, reset_obs) — causing Q-value divergence → NaN. Now: build obs from terminal state first, then reset. maybe_reset calls _reset (no-op) which returns obs from the already-reset state for the next step's input. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- PPO with GAE, compile support, 2-GPU split (env/training) - SAC with auto-tuned alpha, compile support, 2-GPU split - Both support humanoid and humanoid_rich environments Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Collector profiling script with torch.profiler and record_function markers - Humanoid numerical accuracy test (mujoco-torch vs MuJoCo C) - Bug report: physics blowup from degenerate states - TODO: fix humanoid contact filtering (spurious constraints) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- HumanoidRichEnv with 336-dim observations (dm_control-style) - SmoothHumanoidRichEnv with vmap for differentiable physics - SHAC loss module: differentiable actor + value bootstrap + auto-tuned entropy - Direct backprop and SHAC training scripts for humanoid_rich - Launch script for parallel experiments Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Remove the auto_reset flag and its control-flow-based reset logic. Instead, always compute the reset candidate (dx0 + noise) and use torch.where to select between current and reset state based on done mask. This is fully compilable with no graph breaks. The old auto_reset had two problems: 1. if/else broke the torch.compile graph 2. TorchRL's maybe_reset would overwrite terminal obs with post-reset obs, corrupting replay buffers for off-policy algorithms (SAC) The branchless approach works with TorchRL's standard reset flow: _step returns terminal obs, then maybe_reset calls _reset which does the actual per-env reset through the normal path. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…fix/torch-compile', 'perf/make-batch-benchmark', 'feat/humanoid-direct-training', 'feat/humanoid-rich-shac', 'feat/ppo-sac-training' and 'feat/profiling-testing' into combined/humanoid-training-pipeline
Data fields have different numbers of dims (qpos is 2D, xpos is 3D, xmat is 4D). Reshape done_mask to (B, 1, 1, ...) matching each field's ndim so torch.where broadcasts correctly. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Data fields have different numbers of dims (qpos is 2D, xpos is 3D, xmat is 4D). Reshape done_mask to (B, 1, 1, ...) matching each field's ndim so torch.where broadcasts correctly. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Humanoid envs have large observation ranges that cause issues with ReLU activations. Uses TorchRL's ObservationNorm with standard_normal=True, initialized from 1000 random steps. Eval env shares the same stats. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
SAC training loop now checks at 5 stages: 1. Collected batch: obs, action, reward for NaN/Inf/extreme (>1e10) 2. Replay sample: all fields 3. Loss values: NaN/Inf/extreme + full context dump (network outputs, obs/action/reward ranges, alpha state) 4. Gradients: param and grad NaN/Inf after backward 5. Parameters: NaN/Inf after optimizer step Each stage saves a .pt diagnostic file and crashes with a detailed message showing exactly where the bad values originated. base.py _step now checks: - qpos/qvel after physics step for blow-up (>1e10, NaN, Inf) - reward after _compute_reward for extreme values Both dump the full env state (qpos, qvel, action, ctrl) for the offending environment. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
rollout() returns [num_envs, T, obs_dim]. With reduce_dim=0 (default), stats end up with shape [T, obs_dim] instead of [obs_dim]. Use reduce_dim=(0, 1) to reduce over both batch and time dimensions. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
load_state_dict fails because eval env's ObservationNorm has empty [0] buffers before init. Set loc/scale directly instead. Also print timing and range info during init_stats for visibility. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1000 rollouts with compile warmup takes >5 min. 50 rollouts with 8192 envs still gives 50*8192*3 = 1.2M samples — plenty for good mean/std estimates. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
- Remove obs_norm.initialized = True (setting loc/scale already marks it initialized; the property has no setter) - Create vmap and frame_skip loop once in constructor instead of per-call in _step, matching the combined branch approach Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…pipeline Brings ObservationNorm initialized fix to combined branch. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Allows passing torch.compile kwargs (mode, backend, fullgraph, etc.) from the CLI via --compile_mode through to the physics step compilation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
PixelRenderTransform triggers a reset during TransformedEnv construction to discover observation spec, which runs through ObservationNorm. Move loc/scale initialization before TransformedEnv creation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
vmoens
force-pushed
the
feat/ppo-sac-training
branch
from
April 8, 2026 16:09
fd714b0 to
5aa8be8
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Test Plan