# 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

```bash
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)

```bash
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.