[Proposal] Unify differentiable environments and APG training infrastructure #848

Proposal

Make differentiable environments and analytic policy gradient (APG) training first-class infrastructure in EmbodiChain.

Support two environment-provider paths through a shared differentiable training contract:

1. Differentiable or neural physics solvers, including the functional dynamics integration introduced in PR #822.
2. Learned robot world models, fitted from reference-physics transitions and used as differentiable environments for policy optimization.

SuperMimic is an end-to-end functional acceptance case for this infrastructure. Public components should express reusable capabilities; G1, LAFAN1, task rewards, and reference hyperparameters belong to the acceptance task and recipe.

Motivation

APG presents a promising research opportunity for robot learning: policies can use gradients through state evolution instead of relying exclusively on sampled policy-gradient estimates.

EmbodiChain already provides basic APG algorithms, differentiable collection, rollout handling, and training components. The next step is to connect these capabilities to multiple sources of differentiable dynamics with explicit state, gradient, and evaluation contracts.

The intended differentiation is an integrated environment-and-training stack for differentiable robotic RL: reusable providers, reliable gradient lifecycles, online dynamics learning, reproducible experiments, and evaluation in reference physics. Performance and learning benefits must be established experimentally.

This foundation should make it possible to study gradient sources, rollout horizons, critic bootstrapping, model error, and model-update schedules without rebuilding the training infrastructure for each experiment.

Related work: PR #822 and issue #732. Extend the existing checkpoint/resume facilities rather than introducing a separate implementation.

Additional context

Architecture

Both provider paths use DifferentiableVecEnv, DifferentiableCollector, and RolloutKind.DIFFERENTIABLE.

Solver-backed environments retain their functional solver state and physics contracts. World-model environments retain their own model state and task context. NeRD is itself a learned dynamics model; the distinction is the integration path, not whether a provider uses a neural network.

Preserve existing apg behavior and add a generic apg_actor_critic implementation with short-horizon objectives, differentiable continuation weights, terminal-value bootstrapping, and lambda-return critic fitting.

Expose a reusable single-update entry point in the differentiable trainer. The online world-model workflow invokes this entry point rather than maintaining another policy-optimization loop.

Keep the existing trainer / policy / algorithm configuration structure. Select the policy-learning environment through trainer.learning_env; an optional online world-model workflow separately identifies its collection and evaluation physics environments.

Engineering and algorithm ownership

Public entry points and lifecycle changes belong to the infrastructure workstream. Algorithm work proceeds against the agreed contracts.

RL components remain in the RL layer; dynamics models and fitting remain in the model/dynamics-learning layers. Pure numerical operations belong to compute. Backend access stays in simulation adapters. Tasks, robot presets, and assets retain their existing owners.

Align these boundaries with #822. Reuse its data provenance and sequence conventions where applicable, while retaining device-resident replay for online learning.

Required behavior

• Observations and rewards used for policy gradients must come from graph-connected state, not detached native buffers.
• detach_state() cuts the graph without resetting or resampling the task.
• Distinguish hard termination, time-limit truncation, imagination-horizon boundaries, and soft termination probabilities. Soft probabilities must not implicitly trigger Boolean resets.
• Preserve terminal observations for bootstrapping; do not substitute observations from an automatically reset episode.
• Freeze dynamics parameters and normalization statistics during policy updates while preserving gradients with respect to states and actions.
• Update model parameters only after dependent rollout/backward work has completed.
• Capture reference-physics next states before automatic reset. Replay owns stored samples and must not retain mutable simulator views.
• Sequence windows and observation histories must respect episode boundaries. Invalid imagination-start requests must fail explicitly rather than retry indefinitely.
• Keep existing normalizer defaults unchanged; expose the numerical options needed to match the reference recipe.
• Checkpoints preserve policy-learning and dynamics-workflow states, optimizers, normalization, sampling statistics, RNG, compatibility signatures, and progress.
• By default, recovery rebuilds the physics environment and refills replay rather than serializing the entire replay and simulator state. Document this as resumable training, not bitwise trajectory continuation.
• Report real transitions, imagined transitions, model updates, policy updates, and outer iterations separately.

Delivery sequence

1. Shared contracts: agree on provider, rollout, critic-input, and update-boundary semantics; validate the common APG path with small environments.
2. Parallel implementation: engineering delivers data/runtime/adapters; algorithm work delivers APG actor-critic and dynamics fitting, with fixed-input numerical and gradient checks.
3. Online world-model integration: connect reference collection, model fitting, differentiable imagination, APG updates, recovery, and independent physics evaluation.
4. SuperMimic acceptance task: add the G1/LAFAN1 task and reference recipe using the public infrastructure. Fix general gaps in their owning layers.
5. Evidence and documentation: deliver reproducible runs, checkpoints, reports, API documentation, project-context updates, and upstream attribution.

#822 continues independently. This proposal covers interface alignment and gradient regression for that path; its Ant and current-Newton deployment-accuracy qualification remain with the existing workstream. Record the exact solver revision and compatible SDK build used for integration checks.

Acceptance criteria

Infrastructure and gradients

• The same APG update interface works with direct differentiable environments and learned world-model environments without G1-specific branches.
• Replay wraparound, episode boundaries, history padding, pre-reset capture, and sample ownership are covered.
• Actor/critic inputs, truncation bootstrapping, terminal observations, and soft continuation are validated.
• Controlled small systems validate state/action gradients and multi-step gradient propagation.
• Reset, detach, frozen-model updates, and parameter-update ordering preserve the intended computation graph.
• Checkpoint recovery restores all learning states and safely refills replay; incompatible layouts, actions, or assets are rejected before state mutation.
• Existing PPO, GRPO, APG, and differentiable-environment behavior remains covered by regression tests.
• Solver-path checks record any unavailable model/SDK qualification explicitly.

World-model autograd correctness, prediction accuracy, and reference-physics policy performance are separate evidence requirements.

SuperMimic functional acceptance

Use Uni-SuperMimic revision ceafbc863c52b3d613d0e19976794c5a31b9b402 as the implementation reference and vibed revision a1508d42c9d1699a52e8233ab16d08c488441ab9 for comparison.

• Pin the matching G1 asset, joint ordering, actuator parameters, collision settings, 40 LAFAN1 clips, reward, and evaluation conditions.
• Use one NVIDIA GPU with Newton/MuJoCo-Warp, 50 Hz control, and 5 ms physics steps.
• Preserve the reference recipe: 4,096 environments, 1,000 replay steps per environment, UTD 0.4, eight-step model fitting, 32-step imagination, 500 model-only warmup iterations, BF16 network computation, and FP32 state arithmetic.
• Run seeds 0, 1, 2 to a budget of 14,000 outer iterations, saving every 2,000 iterations.
• For each seed, at least one checkpoint within budget reaches a macro normalized return of at least 0.82 and 0/40 failed clips.
• Evaluate each clip from its first frame to completion in independent reference physics, with the training ten-second timeout removed.
• Preserve all checkpoint results and failures; report first qualifying checkpoint time and final-checkpoint performance.

These are acceptance targets to verify, not claims that EmbodiChain has already reproduced the results.

Performance and artifacts

• Record initialization, prefill, compilation, collection, model fitting, policy/critic updates, and evaluation time separately.
• Report Torch allocator measurements alongside native/process GPU-memory coverage and CPU memory.
• Complete the default recipe on the target single GPU without persistent memory growth.
• Deliver run manifests, checkpoints, frozen configs, asset/data signatures, per-clip evaluations, curves, and failure diagnostics.
• Run focused validation, API/context checks, and the repository-required Black 26.3.1 formatting checks before commits.

Default acceptance hardware: one RTX PRO 5000 72GB Blackwell, with the three seeds run sequentially. Establish measured timing rather than assuming the upstream 1.5-hour result transfers across GPU models.

Initial scope boundaries

Multi-GPU training, ROCm/MPS, additional physics backends, real-robot deployment, downstream goal encoders, and a large matched PPO comparison are follow-up work.

References:

• Neural dynamics and solver integration: PR #822
• SuperMimic paper
• Uni-SuperMimic reference
• supermimic_vibed reference

Checklist

• I have checked that there is no similar issue in the repo (required). Related issues were reviewed; no equivalent infrastructure proposal was found.