Skip to content

Humanoid training - #51

Open
vmoens wants to merge 58 commits into
mainfrom
direct-optim-humanoid
Open

Humanoid training#51
vmoens wants to merge 58 commits into
mainfrom
direct-optim-humanoid

Conversation

@vmoens

@vmoens vmoens commented Mar 20, 2026

Copy link
Copy Markdown
Owner
  • Add Humanoid zoo env and direct backprop training script
  • Handle video encoding failure gracefully in direct backprop training
  • Add --compile flag for torch.compile on physics step
  • Compile vmapped step with fullgraph=True, add fwd/bwd timing
  • Add --compile_mode flag to test compile(vmap) vs vmap(compile)

vmoens and others added 30 commits March 20, 2026 14:46
@vmoens @claude
Add Humanoid zoo env and direct backprop training script
69d9152 
Add HumanoidEnv to the zoo (Humanoid-v4 semantics: forward velocity +
healthy reward + ctrl cost).  Add train_direct_humanoid.py that optimizes
a policy by backpropagating through the differentiable mujoco-torch
simulation using differentiable_mode().

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Handle video encoding failure gracefully in direct backprop training
ab7167e 
Wrap wandb.Video in try/except so missing moviepy doesn't crash training.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add --compile flag for torch.compile on physics step
7108b29 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Compile vmapped step with fullgraph=True, add fwd/bwd timing
e2eff33 
- torch.compile the vmap(step) not the inner step
- Use fullgraph=True to verify full graph capture
- Log forward/backward times separately to wandb

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add --compile_mode flag to test compile(vmap) vs vmap(compile)
0add981 
Replaces the boolean --compile flag with --compile_mode choosing between
compile(vmap(fn)), vmap(compile(fn)), or no compile. Keeps --compile as
a deprecated alias. Forward/backward timing already present from prior commit.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Rewrite direct backprop training to use TorchRL envs throughout
b964329 
- SmoothHumanoidEnv subclass with sigmoid-based reward and no hard
  termination for gradient-friendly differentiable training
- Uses env.step() inside differentiable_mode() instead of manual
  vmap/step loop — TorchRL env handles batching internally
- BatchNorm on observations for implicit exploration noise
- Default batch size 1024, horizon 15 (short horizon to limit
  gradient explosion through contacts)
- Gradient accuracy test suite using torch.autograd.gradcheck:
  single step, multi-step, smooth reward, and full env step
  verified against finite differences in float64

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Relax multi-step gradcheck tolerance for accumulated numerical error
a76923a 
3 consecutive physics steps accumulate small floating-point differences
between analytical and finite-difference gradients. Widen eps/atol/rtol
to account for this while still catching real gradient bugs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix BatchNorm device placement — pass device to BatchNorm1d
452e0fe 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add generic train_direct.py supporting all zoo envs
5257135 
New --env flag selects the environment (halfcheetah, humanoid, ant, etc.).
DifferentiableEnv wrapper automatically adds fixed_iterations=True and
smooth reward overrides for envs with hard thresholds. HalfCheetah reward
is already smooth so no override needed.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add --no_fixed_iterations flag for envs where fixed solver diverges
bc9ce70 
HalfCheetah and Ant produce NaN with fixed_iterations=True.
This flag lets us try differentiable mode with the default
variable-iteration solver to see if useful gradients still flow.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add --grad_normalize option for gradient normalization
d2f5a38 
When gradients explode (norms >10^10), standard clipping effectively
zeroes the gradient direction. Normalization preserves the direction
while controlling the step size.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add PPO and SAC training scripts using TorchRL
e0284da 
Standard RL baselines for mujoco-torch zoo environments.
Both use TorchRL losses, collectors, transforms, and replay buffers.
Large batch sizes with vmap-batched envs for fast training.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix TorchRL v0.11 API: critic_coef -> critic_coeff
8ece6ac 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC: use make_value_estimator for gamma (TorchRL v0.11 API)
44161c8 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC ValueEstimators import - use default_value_estimator
ec93375 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC optimizer: extract flat tensor params from TensorDict
6139294 
loss_module.actor_network_params.values() returns TensorDicts.
Need values(True, True) to get individual parameter tensors.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC: pass action_spec to SACLoss for target_entropy inference
73ed543 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC inplace operation error with retain_graph=True
9fa9f06 
SACLoss computes all losses in a single forward pass with shared
computation graphs. Sequential backward calls on individual losses
fail because earlier steps modify tensors needed by later losses.
Adding retain_graph=True on qvalue and actor backward passes fixes this.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC: single combined backward pass to avoid inplace modification
ecd9363 
The previous approach of sequential backward+step caused inplace
modification errors because critic_optim.step() modified Q-network
weights before actor loss could backpropagate through them. Now all
three losses are summed and a single backward() is called before
any optimizer steps.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC alpha logging: use log_alpha.exp() instead of .alpha
6377d93 
SACLoss in TorchRL v0.11 doesn't expose .alpha directly.
Compute it from log_alpha instead.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add HumanoidRichEnv with 336-dim observations
5e40fc8 
The standard humanoid env (53-dim obs: qpos+qvel) lacks critical
proprioceptive data for learning locomotion. HumanoidRichEnv adds
cinert (body inertia), cvel (COM velocity), and qfrc_actuator
(actuator forces), matching Gymnasium Humanoid-v5's observation space.

New files:
- mujoco_torch/zoo/humanoid_rich.py: registered as "humanoid_rich"
- examples/train_direct_humanoid_rich.py: direct backprop variant
- run_humanoid_rich.sh: deploys PPO + SAC + direct on 3 GPUs

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add humanoid numerical accuracy test (mujoco-torch vs MuJoCo C)
3030721 
8 tests comparing dynamics, rewards, action effects, and zoo env
behavior between mujoco-torch and reference MuJoCo C backend.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add TODO: fix humanoid contact filtering (spurious constraints)
20a126c 
mujoco-torch includes non-penetrating contacts in the constraint set
that MuJoCo C filters out, causing divergent humanoid dynamics with
ctrl and blocking locomotion learning.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix solver line search to match MuJoCo C's PrimalSearch algorithm
36462b9 
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>
@vmoens @claude
Add SHAC loss module and training script for humanoid rich
67cd0de 
SHAC combines short-horizon differentiable rollouts with a learned value
function (terminal bootstrap) and stochastic policy (entropy-regularised)
for exploration, addressing the lack of exploration in direct backprop.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SHAC: use action_log_prob key (TorchRL convention)
a0d59e8 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SHAC: detach rollout before critic loss to avoid double backward
9b54f44 
The actor backward consumes the physics computation graph. The critic
loss doesn't need physics gradients, so detach the rollout first.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Use multiplicative reward in SmoothHumanoidRichEnv (dm_control style)
389a981 
Forward velocity reward is now gated by the soft healthy factor, so the
agent only gets locomotion reward while upright. Prevents reward hacking
via falling forward (high velocity, short horizon) and flying exploits.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC: separate backward passes for critic/actor/alpha
0105cdd 
Combined backward caused gradient interference between actor and critic,
leading to Q-value divergence (8e+15). SAC requires separate backward
passes. Also add grad clipping on actor parameters.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix SAC: use retain_graph and defer optimizer steps
01d4705 
TorchRL SACLoss shares the computation graph between losses. Must
backward all losses before stepping any optimizer, using retain_graph
to keep the shared tensors alive.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
vmoens and others added 28 commits March 23, 2026 14:41
@vmoens @claude
Add --compile flag to PPO/SAC, enable adaptive integration for direct…
90c0449 
…/SHAC

- PPO/SAC: add --compile flag to torch.compile the physics step
- PPO: scale up to 8K envs, 8M frames/batch, 2048 mini-batch
- SAC: scale batch_size to 8K, total_frames to 50M
- Direct/SHAC: enable --adaptive_integration for correct contact gradients

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
SAC: collect 8M frames/batch (8K envs × 1000 steps) to match PPO scale
c3ea7c4 
The collector was only doing 1 step per iteration (256 frames), making
compile useless and throughput bottlenecked by gradient updates. Now
collects 8M frames per batch with 8K envs, matching the PPO setup.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Split PPO/SAC across 2 GPUs: env collection on one, training on another
4caea5b 
Add --train_device flag to both scripts. Collector uses storing_device
to transfer data directly to the training GPU. This fixes PPO OOM
(8K envs + 8M batch on single 140GB GPU) by splitting memory across two.

GPU layout: PPO on 0+4, SAC on 1+5, Direct on 2, SHAC on 3.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add degenerate physics state detection in SAC + bug report
69390a3 
SAC now detects extreme rewards (|reward| > 1e6) and saves the 10-step
trajectory slice leading to the blow-up as a .pt file for investigation.
Bug report documents the Cholesky failure and degenerate state issue
that occurs with 2048+ parallel humanoid envs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Update bug report with analysis of captured degenerate state
0e39423 
Single physics step causes cfrc_ext (contact forces) to spike to 1450
from normal ~50-120. Obs indices 256-269 correspond to cfrc_ext in the
rich observation. No NaN/Inf, just extreme finite values from solver.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix compile_step: compile(vmap(fn)) instead of vmap(compile(fn))
ab98d45 
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>
@vmoens @claude
Add throughput benchmark to isolate physics vs env vs collector overhead
cc9c7b8 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix bench: call reset() before accessing _dx
8193f28 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fuse frame_skip into compiled physics + use update_ instead of replace
bf07007 
- 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>
@vmoens @claude
Refactor benchmark to test TorchRL faster-collector branch optimizations
8e040b8 
Adds fast-path flags (_trust_step_output, _skip_maybe_reset, update_traj_ids=False)
and compares standard vs fast collector throughput.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix fast-collector benchmark: don't skip maybe_reset (env needs it)
88ebb07 
MujocoTorchEnv handles resets in _reset() via the collector's auto-reset,
not inside _step(). _skip_maybe_reset requires self-resetting envs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add --collector_only flag to benchmark to skip physics tests
cd433bf 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add collector profiling script with torch.profiler
c6570f6 
Profiles 100 collector steps after warmup, exports Chrome trace + summary tables.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix profiler: disable with_stack and stack grouping to avoid OOM
7c88828 
with_stack=True + group_by_stack_n=3 used 110GB RAM for summary tables.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Reduce profiler overhead: fewer warmup steps, no record_shapes
24a7bca 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Profile with eager physics (no compile) to avoid recompilation stalls
ea7e317 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Reduce profiler scope: 64 envs, 50 steps to keep trace under 1GB
3390e9d 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Add record_function markers to profiler for readable traces
e7fe35a 
Instruments: ENV._step, _physics_step, _build_obs, _compute_reward,
_compute_terminated, _reset, EnvBase.step, step_and_maybe_reset,
maybe_reset, _step_proc_data, TransformedEnv._step, each Transform,
Actor.forward, and Collector methods.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix torch.compile recompilation from variable nefc shapes
5c785f3 
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>
@vmoens @claude
Add --compile flag to profiler to profile compiled physics step
2d0e650 
Now that the recompilation fix (5c785f3) is in, we can profile with
torch.compile enabled to see actual collector overhead when physics
is fast (~16ms vs 739ms eager).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Optimize _make_batch: expand+clone instead of n clones+stack
186c817 
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>
@vmoens @claude
Add auto_reset: fuse env reset into _step to skip TorchRL maybe_reset
ab82959 
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>
@vmoens @claude
Add --auto_reset and --fast_collector flags to SAC training script
13abe3e 
Wire up the fused auto-reset and fast collector optimizations
(trust_step_output, skip_maybe_reset, update_traj_ids=False).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix eval env: don't compile or auto_reset for single-env eval
bf86678 
Compiling for 1 env triggers a separate JIT compilation that wastes
minutes. Eval is infrequent and doesn't need these optimizations.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Fix auto_reset NaN: keep maybe_reset for InitTracker, make _reset no-op
e89bc02 
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>
@vmoens @claude
Fix auto_reset: build obs before resetting to preserve terminal obs
7c79e16 
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>
@vmoens @claude
Fix SmoothHumanoidRichEnv: vmap physics step for multi-env batches
53f6859 
The overridden _physics_step was a raw step() call without vmap,
causing IndexError when num_envs > 1 (batch dim mistaken for joint dim).
Match the base class pattern: vmap for multi-env, raw for single-env,
with frame_skip loop.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@vmoens @claude
Robust Cholesky: raise inline threshold to 32 + clamp pivots
97a1984 
- 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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@vmoens