amsks/bilinear-checkpoints

FB / GC-IQL Bridge — Checkpoints

Model checkpoints from a study of future-reachability representations (Forward–Backward / successor-measure methods, GC-IQL, contrastive RL, PSM/RLDP, TD-MPC2) on OGBench cube manipulation, antmaze, and scene. The collection backs the "bridge" analysis of why successor-measure objectives underperform scalar Bellman supervision on manipulation, and the proposed Bellman-supervised FB and centered-measure variants.

• 2,529 checkpoints (.pt torch, .pkl jax flax params), all training save-intervals.
• ~253 GB, spanning 2 frameworks × many methods × tasks × seeds.
• Every run's original training config travels with it, and a top-level checkpoint_manifest.json indexes everything.

Datasets are not included. These are trained agents, not data. The offline datasets are the public OGBench releases — obtain them from OGBench, not here.

Repository layout

checkpoints/
  <framework>/               # torch | jax
    <method>/                # e.g. fb-flowbc, RGFB-Cube-Single-State-d50-v0, tdmpc2
      <task>__s<seed>__<run>/
        step_200000.pt        (torch)   or   params_500000.pkl   (jax)
        ...                                   ...
        config.yaml           (torch: the run's .hydra/config.yaml)
        flags.json            (jax:   the run's flags.json)
checkpoint_manifest.json      # full index (see below)

Labels (method, task, seed) are read from each run's authoritative config, not parsed from paths.

What's inside

Framework	Files	Notes
torch (.pt)	1,530	FB / FlowBC, PSM-FlowBC, RLDP-FlowBC, CRL-FlowBC, TD-MPC2
jax (.pkl)	999	GC-IQL flow variants: bilinear, centered-measure, RG-FB, HIQL/GCIVL

Representative methods: fb-flowbc, psm-flowbc, rldp-flowbc, tdmpc2, RGFB-Cube-Single-State-d50-v0 (Bellman-supervised FB), GCIQL-State-Flow-CenteredMeasure-{Aux,Only}-d50 (centered successor measure), GCIQL-State-Flow-Bilinear-*, HIQL-AntMaze-Medium-*. Tasks: cube-single-play-v0 (state & pixels), antmaze-medium-navigate-v0, scene-play-v0.

The manifest

checkpoint_manifest.json is a JSON array with one record per checkpoint:

{
  "repo_path": "checkpoints/jax/RGFB-Cube-Single-State-d50-v0/cube-single-play-v0__s0__rgfb_v0_s0/params_500000.pkl",
  "framework": "jax",
  "method": "RGFB-Cube-Single-State-d50-v0",
  "task": "cube-single-play-v0",
  "seed": 0,
  "step": "500000",
  "size": 89070661
}

Downloading

Grab a single checkpoint (only that file is fetched):

from huggingface_hub import hf_hub_download

path = hf_hub_download(
    repo_id="amsks/bilinear-checkpoints",
    filename="checkpoints/jax/RGFB-Cube-Single-State-d50-v0/"
             "cube-single-play-v0__s0__rgfb_v0_s0/params_500000.pkl",
)

Find checkpoints by filtering the manifest:

import json
from huggingface_hub import hf_hub_download

manifest = json.load(open(hf_hub_download("amsks/bilinear-checkpoints",
                                          "checkpoint_manifest.json")))

# e.g. every RG-FB seed
rgfb = [r for r in manifest if r["method"] == "RGFB-Cube-Single-State-d50-v0"]
for r in rgfb:
    ckpt = hf_hub_download("amsks/bilinear-checkpoints", r["repo_path"])

Download an entire method/run folder:

from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="amsks/bilinear-checkpoints",
    allow_patterns="checkpoints/jax/RGFB-Cube-Single-State-d50-v0/*",
)

Loading a checkpoint

Checkpoints are training-state artifacts, not standalone modules — construct the matching agent from the training code, then load the weights. Each run folder ships the exact config used to train it.

• jax (params_*.pkl) — pickled flax parameter pytree for the GC-IQL flow agent. Build the agent from the co-located flags.json, then restore the params.
• torch (.pt) — a torch.load-able state dict for the FB / FlowBC / TD-MPC2 agent. Build the agent from the co-located config.yaml (.hydra), then load_state_dict.

The agent definitions and loaders live in the Factored-FB codebase (tools/wandb_mode_shim/gciql_flow.py for jax GC-IQL; agents/fb/… for torch FB). See that repository for exact construction.

Provenance & reproducibility

• step is the training step of the save; multiple steps per run are included so training trajectories can be reconstructed.
• Configs are byte-for-byte the ones used at train time (config.yaml / flags.json).
• Three identical duplicate backups were de-duplicated; no run is lost.

License

MIT.