JAX/Flax probabilistic dynamics and MBPO experiments; PPO, SAC, TRPO Hessian–vector solvers and REDQ critic ensembles.
Public algorithm implementations and model-based experiments; latent-world-model research documented separately.
Technical deep dive Architecture, implementation & tradeoffs
World models.
Policy optimisation.
World models
Learn dynamics.
Reason over futures.
Probabilistic dynamics and model-based rollouts lead this branch of RLx-Core. JAX/Flax modules predict mean and log-variance; PPO and SAC implementations form part of the broader policy-learning toolkit. Latent-world-model research connects this work to PlaNet and Dreamer.
- 01
Collect
(sₜ, aₜ, rₜ, sₜ₊₁)Real transitions anchor the learning problem. The MBPO scaffold separates real and model replay.
- 02
Fit dynamics
μϕ(s,a), log σ²ϕ(s,a)The implemented Flax module has independent mean and log-variance heads.
- 03
Branch rollouts
sₜ → ŝₜ₊₁ → … → ŝₜ₊ₖThe MBPO design uses short model-generated trajectories; rollout-length ablations are configured.
- 04
Improve policy
real replay + model replayThe intended mixed-data update closes the loop. End-to-end integration remains in progress.
Probabilistic dynamics + policy
Dense JAX/Flax dynamics with separate mean and log-variance outputs. A tanh-squashed Gaussian policy, clipped log standard deviation and explicit PRNG splitting.
Franka manipulation scene
USD scene composition, articulated joints, rigid-body props, collision and friction configuration, surface-gripper control and joint-state input.
World-model research track
From the RLx / RLTK roadmapShort branched rollouts
Generate short synthetic trajectories from real replay states, then use them for policy improvement.
Probabilistic dynamics for planning
Propagate model uncertainty through an ensemble and trajectory sampling when evaluating candidate actions.
Latent-state planning
Encode observations into a recurrent latent state, then evaluate action sequences through learned dynamics.
Learning in imagination
Use a learned world model to improve behaviour through imagined trajectories.
A learned model for search
Learn representations and dynamics that support reward, value and policy prediction during planning.
Visual prediction for policy learning
Train a video-prediction model of the environment, then use the learned simulator to improve a policy.
The public branch contains the dynamics and policy modules, MBPO scaffolding and Isaac Sim scene work. Ensemble training, synthetic-rollout integration and full latent-agent training are further work in this track.
Policy optimisation
The update is
part of the system.
Curvature, target bias and data reuse are explicit in these implementations. The interesting part is what happens between collecting a transition and accepting a new policy.
Clipped probability-ratio objective
Gaussian policy, clipped surrogate, entropy regularisation, minibatch epochs, separate value regression and gradient-norm clipping.
Entropy-regularised actor–critic
Twin Q networks, a stochastic Gaussian actor, replay-buffer updates, a target value network and learned entropy temperature.
Curvature without a dense Hessian
- 01
Differentiate
g = ∇θ L(θ)Flatten autograd derivatives of the policy-ratio surrogate.
- 02
Solve
v → Hₖₗv + λvDifferentiate the gradient–vector product; solve iteratively with residual-based termination.
- 03
Scale
½ ΔθᵀHΔθ ≤ δScale the candidate step using the KL curvature estimate.
- 04
Accept / revert
KL budget ∧ improvementBacktrack until the surrogate improves inside the trust region; restore the old parameters if no candidate passes.
Reuse experience without one critic owning the target
Illustrative subset selection: two critics contribute to the target minimum; the policy objective uses the ensemble mean.
Randomise the target subset, regress the critic ensemble, optimise the policy against the mean Q estimate, learn the entropy temperature, then softly update the target networks.
Runtime contract
A policy export
is a control interface.
The industrial design carries a trained policy into an inference-only runtime with declared shapes, normalisation, action bounds, timing and fallback responsibilities.
- 01
Train
Python / JAX → frozen parametersVersion the experiment, data assumptions and policy checkpoint.
- 02
Package
model.onnx + metadata.jsonBind the exported graph to the observation/action contract and model hash.
- 03
Execute
acquire → normalise → inferA fixed-tick C++20 runtime with preallocated working buffers is the target execution design.
- 04
Constrain
clamp → rate / jerk → actionMonitor validity and timing; signal the classical-control path when takeover is required.
What travels with the model
- abi_version · dimensions
- Interface version, obs_dim and act_dim define compatibility.
- obs_norm · preprocess
- The runtime receives the same normalisation and preprocessing used in training.
- act_bounds · tick_hz
- Action limits and the expected control period are explicit.
- training_commit · onnx_sha256
- The binary resolves to a particular training revision and content hash.
- kpi_ref · license
- Evaluation context and usage terms travel with the artifact.
How the learned policy gets control authority
Shadow
u = uclassicalThe learned policy observes and logs counterfactual actions. The classical controller retains command authority.
Compare proposed actions with the actual execution before handover.Residual
u = guard(uclassical + Δupolicy)The learned policy contributes a bounded correction over the classical controller. The combined command passes through the actuator envelope.
Bound the residual and define the trigger for returning to the classical controller.Primary
u = guard(upolicy)The learned policy supplies the main action within the declared execution envelope. Deadline, novelty and input monitors can request fallback.
Control authority depends on runtime validity and a defined takeover path.Cooperative
u = guard(blend(uclassical, upolicy))A scheduler coordinates the contribution of classical and learned policies. Blending and transitions become explicit parts of the controller design.
The scheduler owns the handover; actuator constraints still apply to the result.Python/C++ parity ≤ 10⁻⁵ L∞; p95 inference ≤ 2 ms on reference hardware; zero post-init allocations; no missed ticks at the target rate; comparison against a classical baseline.
These are design gates for the deployment track, not reported benchmark results.Model-based learning & simulation
- Wrote JAX/Flax dynamics and stochastic-policy modules. The dynamics network returns separate mean and log-variance predictions; the policy uses a tanh-squashed Gaussian with clipped log standard deviation and explicit random-key handling.
- Set up the MBPO architecture around separate real and model replay paths, rollout-length configurations and ablation scripts. The integrated ensemble training and model-rollout loop remain experimental work.
- Built a Franka manipulation scene in Isaac Sim with USD assets, articulated joints, rigid-body props, collision and friction configuration, surface-gripper control and joint-state input.
Second-order policy optimisation
- Implemented the TRPO policy-ratio surrogate, flattened autograd gradients and conjugate-gradient solving with residual-based termination.
- Computed damped KL Hessian–vector products through higher-order autograd, avoiding explicit construction of a dense Hessian.
- Scaled the candidate step to a KL budget, backtracked for surrogate improvement, and restored the previous parameters when no step passed.
- Implemented PPO with Gaussian action log probabilities, minibatched clipped ratios, entropy regularisation, value regression and gradient clipping.
Ensemble value estimation & experience reuse
- Implemented REDQ with 10 critics, a random 2-critic target subset and a configured update-to-data ratio of 20.
- Used the subset minimum for target estimation, the ensemble mean in the policy objective, learned entropy temperature and soft target updates.
- Built TD3, SAC and value-learning implementations spanning Double/Dueling DQN, prioritised replay and Rainbow, alongside tabular planning, Dyna-Q and prioritised sweeping.
From learned dynamics to industrial world models
- The RLx-Core and RLTK roadmaps extend model-based control through MBPO, PETS and PlaNet into Dreamer V1/V2/V3, MuZero and SimPLe.
- The industrial research questions include long-horizon manipulation, partially observed machine state, failure prediction, experience reuse and adaptation across tasks.
- The wider programme includes offline and imitation learning, hierarchical skills, multi-agent coordination, sequence-model policies, robust sim-to-real transfer and risk-sensitive control.
- The deployment design binds a policy to observation normalisation, actuator bounds, model identity, timing assumptions and Shadow, Residual, Primary or Cooperative operating modes.
Prediction becomes useful when it changes a control decision.
A world model needs to retain the dynamics that matter for action: contact, timing, hidden state and how errors compound over a rollout. That connects the research back to controller design. The policy needs a useful imagined future; the runtime still needs an action it can execute, a deadline it can meet and a defined response when its assumptions fail.
Implemented policy and value-learning families; experimental model-based modules and simulation work; world-model and industrial deployment research. The interactive architecture shows the implementation and research layers explicitly.