D4RL MuJoCo Code for ES-SSM

The offline RL entry point runs the D4RL MuJoCo experiment grid:

9 datasets (HalfCheetah / Hopper / Walker2d × medium / medium-replay / medium-expert)
3 seeds per cell
ES-SSM at full capacity K_max=32 (esssm_K32), with embedding dimension 72 and three layers, followed by compact models at K ∈ {2, 4, 8, 16, 32}.

After training, each ES-SSM(K_max=32) cell is sliced into five independent compact checkpoints at K ∈ {2, 4, 8, 16, 32}. Compact inference loads each file in isolation; its on-disk size, runtime memory, and forward-pass latency reflect the deployed K and not the source K_max.

File layout

d4rl_mujoco_RL.py
runtime_utils.py
reproduce_DL.sh
requirements-rl.txt
docs/offline_rl.md

The script keeps three separate working roots, all configurable on the CLI:

--hdf5-dir   ./data
--cache-dir  ./cache_filters
--project-root ./project

Environment

Run the commands below from the repository root. Install the dependencies with python -m pip install -r requirements-rl.txt.

Python ≥ 3.10 with:

torch >= 2.4
numpy
h5py
gymnasium >= 1.0
mujoco >= 3.0

Atari is not used. CUDA is recommended for the full sweep but the script also runs end-to-end on CPU (slowly).

Reproducing the headline grid

python d4rl_mujoco_RL.py download-data \
    --hdf5-dir ./data

python d4rl_mujoco_RL.py sweep-train \
    --project-root ./project --hdf5-dir ./data --cache-dir ./cache_filters \
    --seeds 42,43,44 --num-steps 150000 --slots-per-gpu 3

python d4rl_mujoco_RL.py sweep-export \
    --project-root ./project --cache-dir ./cache_filters \
    --seeds 42,43,44 --Ks 2,4,8,16,32

python d4rl_mujoco_RL.py sweep-eval \
    --project-root ./project --cache-dir ./cache_filters \
    --seeds 42,43,44 --num-eval-episodes 100 --slots-per-gpu 3

python d4rl_mujoco_RL.py tabulate \
    --project-root ./project --seeds 42,43,44 --out-dir ./results

reproduce_DL.sh runs these five commands in sequence with the same defaults.

Running a single cell (debugging)

python d4rl_mujoco_RL.py train \
    --env hopper --dataset medium --model esssm --embed-dim 72 --K-max 32 \
    --seed 42 --num-steps 150000 \
    --hdf5-dir ./data --cache-dir ./cache_filters \
    --out-dir ./project/runs/esssm_K32/hopper-medium

python d4rl_mujoco_RL.py export-compact \
    --ckpt ./project/runs/esssm_K32/hopper-medium/ckpt.pt \
    --out-dir ./project/runs/esssm_K32/hopper-medium/compact \
    --Ks 2,4,8,16,32 --cache-dir ./cache_filters

python d4rl_mujoco_RL.py eval \
    --ckpt ./project/runs/esssm_K32/hopper-medium/compact/esssm_K8.pt \
    --out-json ./project/runs/esssm_K32/hopper-medium/compact/eval_K8.json \
    --num-eval-episodes 100 --cache-dir ./cache_filters

What gets written

For each training cell runs/<tag>/<env>-<dataset>[/seed<N>]/:

ckpt.pt
summary.json
progress.jsonl
train.log
eval.json
compact/

After tabulate:

results/per_seed.csv
results/per_env_mean.csv
results/by_K.csv

Hyperparameters

Archived training settings used by this entry point:

Setting	Value
Embedding dim	72
Layers	3
Context length	20
Batch size	64
Optimizer	AdamW(β=(0.9, 0.999), wd=1e-4)
Learning rate	1e-4 with linear warmup over 10k steps then constant
Total steps	150,000
Gradient clip	0.25
Eval episodes	100 (final)
Eval target return	Higher of (3600, 1800) for Hopper, (12000, 6000) for HalfCheetah, (5000, 2500) for Walker2d

ES-SSM uses budget-dropout: after 10% of training at full K_max=32, K_rt is sampled log-uniformly from {2, 3, 4, 5, 6, 8, 12, 16, 24, 32}; every 8 steps the full K_max is anchored. This is what makes prefix-truncated compact checkpoints reliable at every K.

Compact deployment

build_compact_esssm_from_prefix constructs a fresh RLSequenceModel whose ES-SSM has K_max = K, copies the K-prefix of sigma, phi_f_ri, M, and the selector output rows from the source checkpoint, and re-saves the model as an independent .pt. Loading the compact file does not require the source.