# RLx-Core / Reinforcement Learning Toolkit
Source: https://kaustubhkr.github.io/projects/reinforcement-learning.html

← Flagship work (https://kaustubhkr.github.io/work.html) 03 / Learning & world models 2024–2026

# World models & policy optimisation

Built a toolkit for testing how learned dynamics and policy updates produce continuous-control behaviour.
System RLx-Core / Reinforcement Learning Toolkit
My contribution Learning-system architecture, algorithm implementation & model-based experiments
Access Public implementation

Engineering & evidence 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
System architecture RLx-Core / Reinforcement Learning Toolkit

## World models. Policy optimisation.

World models Policy optimisation Runtime contract

### World models

Experimental work · research direction

### 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 replay The intended mixed-data update closes the loop. End-to-end integration remains in progress.

Modules I wrote

#### 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.

Simulation work

#### 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 roadmap
MBPO Research track

##### Short branched rollouts

Generate short synthetic trajectories from real replay states, then use them for policy improvement.

Reference ↗ (https://arxiv.org/abs/1906.08253)
PETS Research track

##### Probabilistic dynamics for planning

Propagate model uncertainty through an ensemble and trajectory sampling when evaluating candidate actions.

Reference ↗ (https://arxiv.org/abs/1805.12114)
PlaNet Research track

##### Latent-state planning

Encode observations into a recurrent latent state, then evaluate action sequences through learned dynamics.

Reference ↗ (https://arxiv.org/abs/1811.04551)
Dreamer V1 / V2 / V3 Research track

##### Learning in imagination

Use a learned world model to improve behaviour through imagined trajectories.

Reference ↗ (https://danijar.com/project/dreamerv3/)
MuZero Research track

##### A learned model for search

Learn representations and dynamics that support reward, value and policy prediction during planning.

Reference ↗ (https://arxiv.org/abs/1911.08265)
SimPLe Research track

##### Visual prediction for policy learning

Train a video-prediction model of the environment, then use the learned simulator to improve a policy.

Reference ↗ (https://arxiv.org/abs/1903.00374)

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.
Model-based branch ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit/tree/main/phase_06_model_based/1_MBPO) Research roadmap ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit#algorithms--roadmap-phases)

### Policy optimisation

Implemented algorithm core

### 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.

PPO / On-policy

#### Clipped probability-ratio objective

Gaussian policy, clipped surrogate, entropy regularisation, minibatch epochs, separate value regression and gradient-norm clipping.

SAC / Off-policy

#### Entropy-regularised actor–critic

Twin Q networks, a stochastic Gaussian actor, replay-buffer updates, a target value network and learned entropy temperature.

TRPO / Trust region

#### Curvature without a dense Hessian

(H KL + λI) Δθ = g Hessian–vector products + conjugate gradients

- 01

#### Differentiate

g = ∇θ L(θ) Flatten autograd derivatives of the policy-ratio surrogate.

- 02

#### Solve

v → Hₖₗv + λv Differentiate 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 ∧ improvement Backtrack until the surrogate improves inside the trust region; restore the old parameters if no candidate passes.

REDQ / Critic ensemble

#### Reuse experience without one critic owning the target

Q 1 Q 2 Q 3 Q 4 Q 5 Q 6 Q 7 Q 8 Q 9 Q 10
Illustrative subset selection: two critics contribute to the target minimum; the policy objective uses the ensemble mean.
10 Critic networks
2 Target subset
20 Configured UTD

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.

PPO implementation ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit/tree/main/phase_04_pg/4_ppo/PPO) SAC implementation ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit/tree/main/phase_05_actor_critic/2_sac/SAC) TRPO update and line search ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit/tree/main/phase_04_pg/3_trpo/trpo) REDQ implementation ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit/tree/main/phase_05_actor_critic/3_redq/REDQ)

### Runtime contract

Deployment architecture · specified in README

### 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 parameters Version the experiment, data assumptions and policy checkpoint.

- 02

#### Package

model.onnx + metadata.json Bind the exported graph to the observation/action contract and model hash.

- 03

#### Execute

acquire → normalise → infer A fixed-tick C++20 runtime with preallocated working buffers is the target execution design.

- 04

#### Constrain

clamp → rate / jerk → action Monitor 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 Residual Primary Cooperative

### Shadow

u = u classical The 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(u classical + Δu policy ) 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(u policy ) 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(u classical , u policy )) 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.

README acceptance targets 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.
Deployment model contract ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit#deployment-model-contract) Related C++ control implementation ↗ (https://github.com/kaustubhkr/industrial-control-toolkit)

01 / Experimental implementation

## 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.

02 / Implementation

## 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.

03 / Implementation

## 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.

04 / Research direction

## 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.

Engineering note

## 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.

Project context 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.
Source on GitHub ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit) Model-based experiments ↗ (https://github.com/kaustubhkr/reinforcement-learning-toolkit/tree/main/phase_06_model_based/1_MBPO)

Keep exploring

## Control, down to the tick.

↗ (https://kaustubhkr.github.io/projects/control-runtime.html) C++ control runtime

END OF SECTION 8: projects/reinforcement-learning.html