trqam-bc-checkpoints

TRQAM BC checkpoints

Behavior-cloning flow policies for the ten OGBench domains used in Trust Region Q Adjoint Matching (TRQAM). Each checkpoint is the pretrained flow that TRQAM and the flow-based baselines are initialised from.

Paper: Trust Region Q Adjoint Matching (NeurIPS 2026)

Checkpoints

The folder names carry no dataset size. Four domains are trained on datasets larger than the OGBench default, as the table shows. The 10M puzzle-4x4 and cube-triple datasets are the first ten 1M shards of the corresponding 100M datasets.

All checkpoints are trained for 300K steps with seed 10001, 10 flow steps, and a batch of 256. BC uses no reward, so a checkpoint serves every task of its domain. config.json in each folder records the full agent configuration and the environment the run was launched with.

Usage

Each params_300000.pkl holds only the BC flow, at the key TRQAM's loader reads:

{"agent": {"network": {"params": {"modules_actor_slow": ...}}}}

Pass it to TRQAM's main.py as --pretrained_actor_path, together with the domain's network configuration from the table (--agent.actor_hidden_dims, --agent.actor_layer_norm, --horizon_length). TRQAM's --horizon_length defaults to 5, so the navigation domains need --horizon_length=1. The loader copies the flow into every actor slot of the agent.

The weights are plain numpy arrays, so the files load with NumPy 1.x or 2.x and any JAX version. Optimizer state, critics and target networks are not included.

Citation

@inproceedings{dong2026trqam,
    author    = {Yonghoon Dong and Kyungmin Lee and Changyeon Kim and Jaehyuk Kim and Jinwoo Shin},
    title     = {Trust Region Q Adjoint Matching},
    booktitle = {Advances in Neural Information Processing Systems},
    year      = {2026}
}