Causal GPT-RL

GPT-style transformers (Llama) running as RL policies in continuous-control environments.

Both LLM generation and RL interaction are autoregressive:

token           → next token                           (LLM generation)
(state, action) → (next state from env, next action)   (RL rollout)

A token here is a completed pair — a state and the action taken in it. Only the action comes from the model; the state comes back from the environment.

A rollout unrolled — each token pairs the current state with the action for it, that action moves the environment, and the environment returns the next state; meanwhile the model generates the next action from completed state-action pairs, not from the state it just received

Causal GPT-RL policies act stably under their own rollouts — long-horizon control without the drift that has historically kept transformers from being usable as RL agents.

A single autoregressive model drives full-episode rollouts via KV cache — no separate critic, no auxiliary networks at inference. The model carries a value head and computes it on every forward pass, but a rollout does not read it: the action alone carries the loop.

The code (GitHub) repository is the public inference runtime. It loads policy bundles, runs Gymnasium/MuJoCo rollouts, and provides small evaluation helpers.

Supported Environments

Trained with Simple Only

Bundle Ctx Return Norm. Simple Ref. Medium Ref.
hopper-v5 32 3158.35±19.90 81.79±0.52 42.65 ✓ 72.91 ✓
walker2d-v5 32 4192.65±85.03 61.21±1.24 59.51 ✓ 83.26 ✗

Trained with Simple + Medium

Bundle Ctx Return Norm. Simple Ref. Medium Ref.
ant-v5 32 5330.29±1666.34 80.08±24.71 59.99 ✓ 86.54 ✗
halfcheetah-v5 32 5795.90±2747.69 36.69±16.65 43.54 ✗ 74.83 ✗
humanoid-v5 32 7817.46±1254.06 90.74±14.77 63.29 ✓ 81.30 ✓
humanoidstandup-v5 32 239740.89±59256.34 76.78±22.47 48.39 ✓ 89.86 ✗

Trained with Medium Only

Bundle Ctx Return Norm. Simple Ref. Medium Ref.
pusher-v5 32 -29.89±5.57 95.89±4.52 94.19 ✓
swimmer-v5 32 258.32±6.23 72.78±1.75 59.44 ✓

Gymnasium environment IDs: Ant-v5, HalfCheetah-v5, Hopper-v5, Walker2d-v5, Humanoid-v5, HumanoidStandup-v5, Pusher-v5, and Swimmer-v5.
HF bundle subfolders are lowercase and case-sensitive (e.g. humanoidstandup-v5).

Training data is expert-free: expert trajectories are not used for training. hopper-v5 and walker2d-v5 use Minari simple-v0 only; the Simple + Medium table uses both tiers; pusher-v5 and swimmer-v5 use medium-v0 only because Minari does not publish a simple dataset for those environments.

walker2d-v5 was measured on a different simulator release than the other rows. See Reproduction runtime, including why its Simple Ref. and Medium Ref. marks are provisional.

Return and Norm. are mean±std over 100 episodes with seeds 0..99. Ctx is context length. max_steps=1000, and KV cache max length is set to Ctx.

Norm. puts random at 0 and expert at 100. Simple Ref. and Medium Ref. are the normalized means of the Minari simple-v0 and medium-v0 datasets, shown for context and not the normalization baseline; marks a reference the bundle's Norm. exceeds, one it does not. Pusher-v5 and Swimmer-v5 show because Minari publishes no simple-v0 for either, and those two bundles were trained on medium-v0 alone.

Reproducing these numbers

To measure a bundle under that protocol, use examples/deploy/reproduce.py: python -m examples.deploy.reproduce --env-id Ant-v5 --episodes 100. The Quick Start's run_episodes seeds only its first reset, so it cannot express 0..99.

Normalized scores use random=0 and expert=100:

100 * (return - random_ref) / (expert_ref - random_ref)

KV cache retention sweep

Why long-context extrapolation can be stable in RL — two measures leave the same first step, a short arrow to the trained window of 32 and a long one to a 1000-step retention, over a single unbroken bar of tokens that runs on past the shorter one

All bundles above share a context_length of 32 — the model's context window used in training, and not a limit at inference. kv_cache_max_len, how much rollout history is retained, is a load-time knob; the headline scores use kv=32 (1×). The wider sweep below shows that retention is environment-dependent: Humanoid is nearly tied at KV32 and KV128, while Swimmer is far weaker at KV32 than at any other retention. Sweeping retention to 8 (0.25×), 32 (1×), 128 (4×), and 1000 (31×) tokens under the same protocol — 100 episodes, seeds 0..99, max_steps=1000:

Bundle kv=8 (0.25×) kv=32 (1×) kv=128 (4×) kv=1000 (31×)
ant-v5 73.62±27.09 80.08±24.71 82.63±21.18 82.66±21.91
halfcheetah-v5 35.74±17.29 36.69±16.65 34.82±18.25 39.01±19.81
hopper-v5 82.02±0.56 81.79±0.52 82.08±6.74 81.53±6.65
walker2d-v5 60.63±1.03 61.21±1.24 57.07±11.82 57.08±11.98
humanoid-v5 86.92±21.35 90.74±14.77 90.75±13.69 84.92±23.00
humanoidstandup-v5 75.44±25.56 76.78±22.47 74.41±25.00 76.11±23.26
pusher-v5 95.60±4.64 95.89±4.52 95.81±4.52 95.81±4.52
swimmer-v5 80.97±2.68 72.78±1.75 80.35±2.90 83.61±1.75

Scores are normalized with random=0 and expert=100. Values are mean±std over 100 episodes.

The kv=32 column repeats the main table's Norm. column; the others are the same protocol at a different retention. hopper-v5's kv=1000 column combines four 25-row batches covering seeds 0..99; its other columns are one 100-row batch.

KV retention interpretation

At kv=128 the rollout attends well past the model's 32-token training window, and kv=1000 runs 31× beyond it — both columns are extrapolation. The 0.25×/1×/4×/31× spacing makes environment-specific effects visible without treating retention as uniformly beneficial.

  • Ant-v5: mean rises with retention and peaks at KV1000; KV128 has the best horizon count in this 100-seed batch (84/100).
  • HalfCheetah-v5: KV1000 has the highest mean, but all four settings have broad return distributions.
  • Hopper-v5: means are flat — 0.55 points separate the best and worst retention. What retention changes here is dispersion: KV8 and KV32 reach the horizon on all 100 seeds, while KV128 and KV1000 end early twice and once, which is where their wider spread comes from.
  • Walker2d-v5: the simple-only step-28,800 bundle is strongest at KV32; KV8 and KV32 complete 100/100 episodes, KV128 completes 93/100, and KV1000 completes 94/100.
  • Humanoid-v5: KV32 and KV128 tie on mean; KV128 is modestly steadier, while KV32 reaches the horizon more often (97/100 versus 95/100).
  • HumanoidStandup-v5: KV32 has the highest mean and lowest dispersion.
  • Pusher-v5: flat across every retention length.
  • Swimmer-v5: KV8, KV128 and KV1000 are all far stronger than KV32.

Reproduction runtime

Every result above can be re-evaluated with the following reference stack:

causal-gpt-rl 0.16.0
torch 2.8.0+cu129
gymnasium 1.2.3
mujoco 3.2.3
minari 0.5.3

The walker2d-v5 measurements are an explicitly documented runtime exception: causal-gpt-rl 0.17.0, torch 2.11.0+cu126, gymnasium 1.1.0, and mujoco 3.8.1. Its KV32 headline is one 100-row batch over seeds 0..99; KV8 and KV128 each combine two 50-row batches, and KV1000 combines four 25-row batches, all covering seeds 0..99. Closed-loop returns depend on runtime and batch width, so compare that row with the reference-stack rows accordingly. And because Simple Ref. and Medium Ref. are derived from trajectories recorded on mujoco 3.2.3, the ✓/✗ marks on that row are provisional — they hold a 3.8.1 measurement against a 3.2.3 reference.

mujoco is pinned to 3.2.3 because that is the version the Minari datasets were recorded with (requirements: ['mujoco==3.2.3', 'gymnasium>=1.0.0']). The Norm. and Medium Ref. columns are derived from those recorded trajectories, so returns are only comparable to them when measured on the same physics.

Install

For Hub loading and MuJoCo environments:

pip install "causal-gpt-rl[hub,mujoco]"

For local development:

git clone https://github.com/ccnets-team/causal-gpt-rl.git
cd causal-gpt-rl
python -m pip install -e ".[hub,mujoco]"

For private bundles, authenticate first:

hf auth login

Quick Start

import gymnasium as gym

from causal_gpt_rl.inference import load_runner_from_hub, run_episodes

env = gym.make("Ant-v5")
runner = load_runner_from_hub(
    repo_id="ccnets/causal-gpt-rl",
    subfolder="ant-v5",
)

stats = run_episodes(env, runner, num_episodes=5, seed=0)
env.close()
print(stats["return_mean"], stats["return_std"])

Notebook version: examples/hub_quickstart.ipynb

Bundle Format

Public bundles use bundle_format_version=2:

bundle/
  model.safetensors
  config.json
  • model.safetensors — model state dict for inference, with state normalization statistics embedded in the weights.
  • config.json — model config, observation specs, action specs, context length, a state_normalization block, and optional env_id.

For a local bundle directory, use load_runner("path/to/bundle").

API

from causal_gpt_rl.inference import (
    PolicyRunner,                          # step-wise rollout policy with KV cache
    load_runner,                           # load runner from a local bundle directory
    load_runner_from_hub,                  # load runner from a Hugging Face Hub repo
    run_episodes,                          # evaluate over N episodes; returns stats dict
    export_bundle,                         # write a bundle directory from a runner
    convert_legacy_bundle_to_safetensors,  # migrate legacy bundles to the safetensors format
)

License

Released under PolyForm Noncommercial License 1.0.0. See LICENSE for details. For commercial licensing, contact the maintainers via ccnets.org.

Downloads last month
423
Video Preview
loading