This repository provides an efficient implementation and evaluation framework for the offline lower confidence bound - value iteration algorithm proposed by Gen Li et al. in "Settling the sample complexity of model-based offline reinforcement learning" (The Annals of Statistics 52.1, 2024). DOI: 10.1214/23-AOS2342
The algorithm performs value iteration with a pessimistic Bernstein-style penalty for uncertain state-action pairs in offline datasets while claiming optimal sample complexity. A key challenge is determining the hyperparameter cb (the penalty scaling factor). While the original paper used a fixed value, performance fluctuates significantly when the same value is applied across different datasets.
To address this, a heuristic is used, that selects the largest cb value that maintains performance sufficiently close to the optimal average cross-validation score on the training data (within a specified margin). The cross-validation uses the doubly robust offline policy evaluation algorithm by Jiang and Li ("Doubly Robust Off-Policy Value Evaluation for Reinforcement Learning").
The heuristic for selecting the right cb hyperparameter is located in src/lcb_vi/eval/grid.py.
The efficient implementation of the pessimistic value iteration is located in src/lcb_vi/solver.py.
The evaluation uses a custom "FrozenLake" environment (modified from the Gymnasium package) with stochastic transitions that introduce risk of falling into holes even under the optimal policy.
Setup. The environment consists of a 4×4 grid where the agent starts at a random field. At each position, the agent can choose from 4 movement directions (up, down, left, right).
Stochastic dynamics. Due to the slippery ice surface, a selected action is executed correctly with 70% probability. However, there is a 30% chance of slipping: each of the three other directions has a 10% probability of being followed instead. If the agent falls into a hole, the episode terminates immediately.
Reward structure. The agent receives a small decaying reward for each step that avoids holes, while the largest reward is obtained by reaching the single goal state as quickly as possible. This creates a trade-off between accumulating steady low rewards versus pursuing potentially higher rewards with greater risk.
Offline datasets are created and managed using the Minari package. The bench command performs comprehensive evaluations across multiple dataset sizes (as specified in the config file), including:
- Grid search for optimal
cbvalues - Full evaluation using actual true environment rollouts
- Performance comparison against naive value iteration and the true policy
- Averaging the findings over multiple runs, which also provides a variance band
Results demonstrate that despite requiring heuristic cb selection, the pessimistic value iteration significantly outperforms naive value iteration, especially on smaller datasets.
- Gym Env:
CustomFrozenLake(modified rewards + transition noise) - Models: Exact tabular
(P, R)from environment and empirical(P̂, R̂)from Minari datasets - Solvers: Standard Value Iteration (VI) and pessimistic variant (
lcb_vi) - CLI: Run experiments, save artifacts/plots, and print metrics tables
uv sync --group dev # create venv and install (incl. dev deps)
uv pip install -e . # or for non-editable environmentuv run python -m projects.frozenlake.genDatasets are stored by Minari in your user data directory and referenced by dataset_id.
List them with:
uv run minari list localThe generator projects.frozenlake.gen creates five Minari datasets per map size, each collected under a different behavior policy to simulate more realistic non-optimal behaviour policies:
- Naming scheme:
custom_frozenlake/<map>_true<P>-v1<map>is the grid shape, e.g.,4x4.<P>∈{0, 20, 40, 60, 80}is the optimal policy override percentage used during data collection.
- Interpretation:
true0→ 0% optimal policy override: all actions are chosen uniformly at random.true80→ 80% optimal policy override: at each step, there is an 80% chance the action is drawn for certain from the optimal policy and a 20% chance the action is chosen uniformly at random.
Examples
custom_frozenlake/4x4_true0-v1— 4×4 map; data collected entirely under the uniform random policy.custom_frozenlake/4x4_true80-v1— 4×4 map; each action has an 80% probability of being surely drawn by the optimal policy.
These datasets differ only in the behavior policy used for collection; environment dynamics and reward definition are identical.
projects/frozenlake/runs/<exp_name>/<dataset_tag>/<YYYY-mm-dd-HHMMSS>__seed<seed>/
├── best_cbs.json # best pessimism scale (cb) per (run, episodes)
├── config.toml # run configuration (as TOML)
├── metrics.json # evaluation metrics
├── run.log # chronological log (info/debug)
├── plots/ # plots saved by `lcb-vi plot`
└── policies/
└── episodes_<ne>/run_<r>/
├── model.npz # P̂, R̂, counts (+ JSON metadata)
├── naive_agent.npz # policy, V, Q (+ metadata)
└── pess_agent.npz # policy, V, Q (+ cb, metadata)
<dataset_tag> is a safe filename produced from dataset_id, e.g.,
custom_frozenlake/4x4_true20-v1 → custom_frozenlake__4x4_true20-v1.
Edit a config file like projects/frozenlake/configs/base.toml:
[eval]
seed = 123
processors = 8 # workers (MP pool)
runs = 10 # evaluation runs per dataset size
[agent]
gamma = 0.99 # discount ∈ (0, 1]
cb = 0.05 # pessimism scale ≥ 0
[data]
dataset_ids = [
"custom_frozenlake/4x4_true0-v1",
"custom_frozenlake/4x4_true20-v1",
"custom_frozenlake/4x4_true40-v1",
"custom_frozenlake/4x4_true60-v1",
"custom_frozenlake/4x4_true80-v1",
]
out_root = "projects/frozenlake/runs"
exp_name = "frozenlake-bench"
# dataset sizes examined (episodes_min..max step)
episodes_min = 5
episodes_max = 2000
episodes_step = 2# 0) generate example datasets
uv run python -m projects.frozenlake.gen
# 1) run the benchmark for all dataset_ids from the config
uv run lcb-vi bench projects/frozenlake/configs/base.toml --processes 8
# 2) make plots (saved next to metrics.json under run_dir/plots/)
uv run lcb-vi plot projects/frozenlake/runs/frozenlake-bench
# 3) print a flattened metrics table from the latest run per dataset
uv run lcb-vi metrics projects/frozenlake/runs/frozenlake-bench --head 20| Command | Description |
|---|---|
uv run lcb-vi bench <cfg_path> [--processes N] [--out-root DIR] |
Run experiments for all data.dataset_ids from the TOML config; writes metrics + artifacts. |
uv run lcb-vi plot <bench_root> |
For each run dir under <bench_root>, render & save plots/eval_panels.(png,svg). |
uv run lcb-vi metrics <bench_root> [--head K] |
Print a flattened table from metrics.json (latest run per dataset). |
bench_rootis by defaultprojects/frozenlake/runs/<exp_name>.
Full Evaluation benchmark of the CustomFrozenLake – true60 environment. All metrics were averaged over 20 runs.
Context. The figure compares naive VI against pessimistic lower-confidence-bound (LCB) VI, evaluated against the true policy from the exact model. The x-axis shows dataset size (episodes used for the empirical model).
Value convergence with expected bias. The “Difference to true Values” panel shows that both methods converge sharply as episodes increase. The larger discrepancy for LCB VI is expected due to the pessimistic bonus scaling, which depends heavily on the heuristically found confidence bound parameter. Despite this, LCB VI converges to true values as data grows.
Better performance with limited data. The pessimistic policy achieves higher average returns than naive VI when data is scarce, steering away from risky transitions as confirmed by lower fall probabilities.
Improved stability. Naive VI shows long trajectories with high variance early on, while pessimistic VI is shorter and more stable. Both converge as data increases.
- The CLI auto-registers
CustomFrozenLakewhen adataset_idstarts withcustom_frozenlake/. - Plots: The 4-panel figure shows (i) value-function error vs. true (optimal), (ii) average return, (iii) probability of falling, and (iv) trajectory length, each as a function of dataset size with uncertainty bands.
MIT — see LICENSE.