Official implementation of ASTEROID (Asymmetric Student–Teacher Rollout In-context Distillation), a simple framework for learning in-context exploration in partially observed decision-making tasks from a privileged expert.
Sequential decision-making under partial observability requires an agent to explore to infer hidden, task-relevant information before acting. Learning such behavior tabula rasa with RL is sample-inefficient, while offline imitation of a clairvoyant expert never demonstrates exploration. ASTEROID resolves this via iterative on-policy asymmetric distillation: the student collects on-policy context under partial observability, and the privileged expert labels actions conditioned on that context. As the context grows, the student learns to infer the hidden task from its own history and act accordingly — recovering Bayesian-posterior-sampling-like exploration while keeping the simplicity of supervised learning.
In-context exploration learned by ASTEROID at test time (vs. the offline DPT baseline):
Key→Door — the first episode searches for the key; every later episode retrieves its location from in-context memory and explores only for what it hasn't seen yet.
Left: Darkroom — cumulative cell visitation over test-time episodes. DPT re-searches the same cells diffusely; ASTEROID sweeps the grid systematically, locks onto the goal, and jumps to expert return. Right: 2D-Navigation — ASTEROID fans hypotheses across the arena and converges on the hidden goal; DPT wanders near the start and never finds it.
ASTEROID trains a history-conditioned Decision Transformer by alternating two phases per iteration (Algorithm 1 in the paper):
- Data generation. Roll out the current student under partial observability to collect on-policy histories, then query the clairvoyant expert for the optimal action at each state. The rollout horizon grows by one episode per iteration (a context-length curriculum).
- Supervised learning. Train the student by conditional imitation of the expert labels on the accumulated dataset.
This on-policy relabeling gives coverage over exactly the histories the student visits at deployment, so it learns to explore without the expert ever demonstrating exploration.
The project uses uv. The core environment (ASTEROID + the DPT and
AAWR baselines) is a single command:
uv syncOptional benchmarks/baselines are installed as extras:
uv sync --extra baselines # RL2 (recurrent-PPO meta-RL)
uv sync --extra procgen # Procgen maze benchmark
uv sync --all-extras # everythingAll commands below can be prefixed with uv run (e.g. uv run python experiments/train_asteroid.py ...).
configs/ per-environment hyperparameter presets (YAML)
datasets/ trajectory dataset + on-policy data collection
environments/ gridworld/navigation/procgen envs, factory, rollout policies
models/ Decision Transformer (+ CNN), DPT transformer, asymmetric critic
experiments/ training and evaluation launchers
scripts/ sweep launchers and plotting helpers
results/ all checkpoints, evaluations, and logs (git-ignored)
Environments: darkroom-easy, darkroom-hard, keydoor-markovian, keydoor-nonmarkovian,
navigation-episodic, navigation-nonepisodic, and the Procgen maze.
Run with a per-environment config preset (any flag overrides the config):
uv run python experiments/train_asteroid.py --config configs/darkroom-easy.yaml --seed 0Or specify arguments directly:
uv run python experiments/train_asteroid.py \
--env_name keydoor-nonmarkovian \
--dagger_steps 5 \
--dataset_size 10000 \
--num_epochs 100 \
--eval_episodes 40 \
--seed 0 --log_wandbSweep over environments and seeds:
ENVS="darkroom-easy darkroom-hard" SEEDS="0 1 2" bash scripts/train_asteroid.shEvery method writes to results/<method>/<exp_name>-<env>-seed<seed>/ with a common layout:
results/asteroid/asteroid-darkroom-easy-seed0/
├── model_args.pkl # model config
├── final_model.pth # final checkpoint
├── dagger_step_<k>/ # per-iteration checkpoints + train/test data
│ ├── model_epoch_*.pth
│ └── eval/ # eval_returns.npz, eval_returns.png, eval_trajs.pkl
Evaluate a trained checkpoint on held-out tasks:
uv run python experiments/eval_policy.py \
--env_name darkroom-easy \
--checkpoint_path results/asteroid/asteroid-darkroom-easy-seed0/dagger_step_4 \
--plot_returnsAll baselines share the same environments, dataset, and results/ layout as ASTEROID.
Trains a Decision Transformer on offline random-context interactions labelled by the expert (no on-policy iterations).
uv run python experiments/train_dpt.py --config configs/darkroom-easy.yaml --seed 0Learns a privileged critic (IQL) then extracts a policy via advantage-weighted imitation.
uv run python experiments/train_aawr.py --config configs/darkroom-easy.yaml --seed 0uv sync --extra baselines
uv run python experiments/train_rl2.py --env_name darkroom-easy --seed 0Sweep the gridworld baselines together:
METHODS="dpt aawr rl2" ENVS="darkroom-easy" SEEDS="0 1 2" bash scripts/run_baselines.shuv sync --extra procgen
# ASTEROID (CNN Decision Transformer) on procgen mazes
uv run python experiments/train_asteroid_procgen.py --seed 0
# Baselines: privileged expert eval, DPT eval (from checkpoint), BC+PPO
uv run python experiments/eval_procgen_expert.py --seed 0
uv run python experiments/eval_procgen_dpt.py --checkpoint <path.pt> --seed 0
uv run python experiments/train_procgen_bc_ppo.py --seed 0
# or the full sweep:
SEEDS="0 1 2" bash scripts/run_procgen.sh@article{poddar2026asteroid,
title = {Asymmetric On-Policy Distillation Learns In-Context Exploration},
author = {Poddar, Sriyash and Bao, Yanda and Krantz, Jacob and Chang, Matthew and
Puig, Xavier and Mottaghi, Roozbeh and Jaques, Natasha and Gupta, Abhishek},
journal = {arXiv preprint arXiv:PLACEHOLDER},
year = {2026}
}