Refresh artefacts and code for the second review release (part 2)
c92019b verified about 2 months ago
35.1 kB
"""Upload the trained Craftax ReMDM checkpoints and results to the Hugging Face Hub.
Discovers every Orbax checkpoint under ``checkpoints/``, every ablation run
under ``experiments/rl_finetuning/outputs/`` and every ``--mode inference``
result under ``results/inference/``, the manuscript figure PDFs under
``results/paper_figures/``, stages them with the repo-relative layout
preserved, drops wandb environment metadata (which carries the author's email,
hostname and local paths), generates a model card from the checkpoints' own
config snapshots, and uploads.
    HF_TOKEN=hf_xxx uv run python scripts/hf_upload.py \\
        --repo-id <ANON_HF_REPO_ID> \\
        [--inference-results PATH ...] [--dry-run] [--private]
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import sys
import tempfile
from pathlib import Path

import yaml

# scripts/ is not a package, so the helper is imported by bare name.
# Python already puts this file's directory on sys.path when the script is
# run directly; the explicit insert is for the file-location loaders the
# tests use (tests/test_config.py), which do not.
_SCRIPTS = Path(__file__).resolve().parent
if str(_SCRIPTS) not in sys.path:
    sys.path.insert(0, str(_SCRIPTS))

from _git_provenance import copy_tracked_file  # noqa: E402

ROOT = Path(__file__).resolve().parent.parent
CKPTS = ROOT / "checkpoints"
RUNS = ROOT / "experiments" / "rl_finetuning" / "outputs"
INFERENCE = ROOT / "results" / "inference"
PAPER_FIGURES = ROOT / "results" / "paper_figures"

PAPER = (
    "Return-Weighted ELBO Fine-Tuning Degrades "
    "Masked Diffusion Planners"
)
CODE_URL = "https://github.com/ANONYMOUS/remdm-planners"
ENV_NAME = "Craftax"

ROLES = {
    "offline": "Diffusion planner (offline BC)",
    "online": "Diffusion planner (online DAgger)",
    "ppo_agents": "PPO-RNN expert",
}

# An ablation run is published as its own summary plus tables and figures; the
# raw per-iteration logs stay in the code repository.
RUN_FILES = ("results.json", "diagnosis.md")
# gdelta/ holds the per-seed gradient measurements behind the paper's
# decomposition appendix, which the checklist promises are openly available.
RUN_DIRS = ("tables", "figures", "gdelta")

# Environment provenance, never needed to restore a checkpoint, and dropped
# from every published config in both repos: the nested `_wandb` blob (email,
# host, git remote, absolute paths), every key containing `wandb_`, every
# `hub_*` key, and `use_wandb` itself — which matches neither prefix, so both
# repos shipped it in the released config while claiming to scrub W&B settings.
# Keys are compared lower-cased, because craftax records them UPPERCASE and
# minihack lower-case.
#
# `wandb_` is a SUBSTRING rule, not a prefix rule, because a W&B key is not
# always at the front: this repo published `RESUME_WANDB_RUN_ID`, and the
# sibling published `baselines_wandb_project` with a real value in two released
# config.yaml files. A suffix rule would have caught the first and left the
# second. The trailing underscore is what makes a substring safe -- `wandbish`
# contains `wandb` but not `wandb_`, so the negative assertions in both repos'
# test_config stand unchanged.
#
# `hub_` stays a PREFIX rule deliberately: as a substring it would swallow
# `github_url`.
DROP_PREFIXES = ("hub_",)
DROP_SUBSTRINGS = ("wandb_",)
DROP_KEYS = ("_wandb", "use_wandb")


def is_environment_key(key: str) -> bool:
    """True for a config key that is provenance rather than recipe."""
    lowered = key.lower()
    return (
        lowered in DROP_KEYS
        or lowered.startswith(DROP_PREFIXES)
        or any(mark in lowered for mark in DROP_SUBSTRINGS)
    )


# wandb-metadata.json is pure environment provenance (email, host, git remote,
# absolute paths) and is never needed to restore a checkpoint.
COPY_IGNORE = shutil.ignore_patterns(
    ".DS_Store", "__pycache__", "*.pyc", "wandb-metadata.json",
)
HUB_IGNORE = ["**/.DS_Store", "**/__pycache__/**", "**/wandb-metadata.json"]


# =============================================================================
# Discovery
# =============================================================================

# `checkpoints/hf/` is where a Hub *download* lands. Publishing from it would
# re-upload already-published artefacts into a nested `checkpoints/hf/...` tree
# on the Hub, so it is never a publish source in either repo.
HF_DOWNLOAD_DIR = "hf"


def _is_download_copy(path: Path) -> bool:
    """True for anything under ``checkpoints/hf/``, wherever it sits."""
    return HF_DOWNLOAD_DIR in path.relative_to(CKPTS).parts


def discover_checkpoints() -> dict[Path, list[int]]:
    """Map each checkpoint directory to its saved step numbers.
    Discovery is at the **released layout**, ``checkpoints/<role>/<name>/<step>/``
    — the layout the Hub repo mirrors, which is why `--dry-run` shows the tree a
    publish would create. A training run writes elsewhere, so its `policies`
    directory has to be copied into place first; the README documents that and
    it is a real requirement, not an accident of this glob.
    Anything under ``checkpoints/hf/`` is skipped as a download copy. The fixed
    depth already excluded it here, one level deeper than a real checkpoint, but
    only by arithmetic — the exclusion is now stated, so it survives a layout
    with a different depth.
    Measured on this repo's live tree: 4 checkpoints discovered, 4 download
    copies under ``checkpoints/hf/`` skipped.
    Returns:
        ``{checkpoint directory: [step numbers]}``.
    """
    models: dict[Path, list[int]] = {}
    for marker in sorted(CKPTS.glob("*/*/*/_CHECKPOINT_METADATA")):
        if _is_download_copy(marker):
            continue
        models.setdefault(marker.parent.parent, []).append(int(marker.parent.name))
    return models


def discover_runs() -> list[Path]:
    """Every ablation output directory holding a ``results.json``."""
    if not RUNS.is_dir():
        return []
    return sorted(d for d in RUNS.iterdir() if (d / "results.json").is_file())


def discover_paper_figures() -> list[Path]:
    """The manuscript figure PDFs built by ``scripts/paper_figures.py``.
    Unlike everything else published here these are *cross-environment*: the
    script reads both this repository's and the MiniHack sibling's
    ``results.json`` and draws the two side by side, so the same PDFs belong in
    both releases and neither repository can build them alone.
    """
    if not PAPER_FIGURES.is_dir():
        return []
    return sorted(PAPER_FIGURES.glob("*.pdf"))


def discover_inference(extra: list[str]) -> list[Path]:
    """Inference result JSONs: the default directory plus any given paths."""
    found: list[Path] = []
    for source in [INFERENCE, *(Path(p) for p in extra)]:
        if source.is_dir():
            found.extend(sorted(source.glob("*.json")))
        elif source.is_file():
            found.append(source)
        elif source != INFERENCE:
            print(f"No inference results at {source}.", file=sys.stderr)
    return list(dict.fromkeys(found))


# =============================================================================
# Helpers
# =============================================================================

def dir_size_mb(path: Path) -> float:
    if path.is_file():
        return path.stat().st_size / 1_048_576
    return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1_048_576


def human_size(path: Path) -> str:
    mb = dir_size_mb(path)
    return f"{mb:.0f} MB" if mb >= 1 else f"{max(mb * 1024, 1):.0f} KB"


def plural(n: int, word: str) -> str:
    return f"{n} {word}" if n == 1 else f"{n} {word}s"


def shorten_paths(value):
    """Shorten absolute cluster paths anywhere in a staged JSON document."""
    if isinstance(value, dict):
        return {k: shorten_paths(v) for k, v in value.items()}
    if isinstance(value, list):
        return [shorten_paths(v) for v in value]
    if isinstance(value, str) and value.startswith("/"):
        return "/".join(Path(value).parts[-2:])
    return value


def environment_key_paths(value, path: str = "") -> list[str]:
    """Every dotted path at which an environment key appears, at any depth.
    The reporting counterpart to :func:`drop_environment_keys`, so a scrub can
    name what it removed rather than claiming a scrub happened. Recursion stops
    at a key that is itself dropped: its whole subtree goes.
    """
    found: list[str] = []
    if isinstance(value, dict):
        for key, sub in value.items():
            name = str(key)
            here = f"{path}.{name}" if path else name
            if is_environment_key(name):
                found.append(here)
            else:
                found.extend(environment_key_paths(sub, here))
    elif isinstance(value, (list, tuple)):
        for i, sub in enumerate(value):
            found.extend(environment_key_paths(sub, f"{path}[{i}]"))
    return found


def drop_environment_keys(value):
    """Strip environment keys at **every** depth, not just the top level.
    Filtering only the top level was a live defect here: `scrub_abs_paths`
    recursed with `shorten_paths` but filtered keys at the top of the document
    only, so `USE_WANDB`, `WANDB_PROJECT`, `WANDB_ENTITY` and
    `WANDB_DOWNLOAD_DIR` survived one level down inside `config_snapshot` and
    were published in both released checkpoints' `resume_metadata.json`, while
    the uploader printed a successful scrub. No credential was exposed -- those
    live in the `_wandb` blob and `wandb-metadata.json`, both already removed --
    but the published surface advertised a W&B account and project that are
    nothing to do with the recipe, which the scrub exists to prevent.
    Mapping types are preserved, so an ordered mapping stays ordered and a
    published structure is unchanged beyond the removed keys.
    """
    if isinstance(value, dict):
        kept = {
            key: drop_environment_keys(sub)
            for key, sub in value.items()
            if not is_environment_key(str(key))
        }
        if type(value) is dict:
            return kept
        try:
            return type(value)(kept)
        except TypeError:
            # A mapping needing constructor arguments (e.g. a defaultdict
            # factory). Structure matters less than not shipping the key.
            return kept
    if isinstance(value, list):
        return [drop_environment_keys(sub) for sub in value]
    return value


def scrub(cfg):
    """Drop the environment keys and shorten absolute cluster paths.
    Both passes recurse. They used to disagree -- ``shorten_paths`` descended
    and the key filter did not -- which is the asymmetry that published W&B
    settings out of a nested `config_snapshot`. Composing them is what stops
    them disagreeing about depth again; special-casing the one nesting we know
    about is what left the general case broken in the first place.
    A document with no environment key anywhere and no absolute path comes back
    equal to what went in, so a caller may skip rewriting it.
    """
    return drop_environment_keys(shorten_paths(cfg))


# Anchored with a lookbehind so a match cannot START at a slash in the middle
# of a RELATIVE path. Without it, `experiments/rl_finetuning/analysis/tables.py`
# matched from the slash after `experiments` and came out
# `experimentsanalysis/tables.py` -- a scrub corrupting a file it had no
# business touching. The character class keeps `@` and `+`: dropping them does
# not merely miss a path, it makes the lookbehind reject the whole match, so an
# absolute path containing either would sail through unshortened.
ABS_PATH_IN_TEXT = re.compile(r"(?<![\w.@+\-])/(?:[\w.@+\-]+/){2,}[\w.@+\-]+")


def shorten_text_paths(text: str) -> str:
    """Shorten absolute cluster paths inside a plain-text file.
    :func:`shorten_paths` only reaches paths that sit in a JSON or YAML *value*.
    Orbax writes its own marker files, and `commit_success.txt` records the
    absolute directory it committed to -- which on this cluster is inside the
    W&B run directory, so each published marker carried the account name, the
    full home path and the run id. That is the same `wandb_run_id` the sidecar
    scrub takes care to remove, published verbatim two directories away.
    Same rule as the structured shortener: keep the last two components.
    """
    return ABS_PATH_IN_TEXT.sub(
        lambda m: "/".join(Path(m.group(0)).parts[-2:]), text
    )


def scrub_staged_json(path: Path) -> list[str]:
    """Scrub a staged JSON file in place; return the paths it dropped."""
    try:
        payload = json.loads(path.read_text())
    except (json.JSONDecodeError, UnicodeDecodeError):
        return []
    dropped = environment_key_paths(payload)
    cleaned = scrub(payload)
    if cleaned != payload:
        path.write_text(json.dumps(cleaned, indent=2))
    return dropped


def scrub_staged_tree(target: Path) -> None:
    """Scrub every JSON and Orbax marker anywhere under a staged directory.
    Staging copied whole trees and scrubbed only the files it knew by name, so
    provenance rode out in the ones it did not: `wandb-summary.json` beside a
    PPO checkpoint, an ablation run's `results.json` (which carried the W&B
    entity into a published release), and Orbax's own `commit_success.txt`.
    Walking the staged tree is what stops the list of known filenames from
    being the thing correctness depends on.
    """
    for f in sorted(target.rglob("*")):
        if not f.is_file():
            continue
        if f.suffix == ".json":
            for key in scrub_staged_json(f):
                print(f"  scrubbed {key} from {f.name}")
        elif f.name.endswith(".txt"):
            text = f.read_text(errors="ignore")
            shortened = shorten_text_paths(text)
            if shortened != text:
                f.write_text(shortened)
                print(f"  shortened cluster paths in {f.name}")


def strip_wandb_block(config_yaml: Path) -> None:
    """Drop the environment keys from a staged config, at every depth.
    Was `_wandb` alone, which left `USE_WANDB`, `WANDB_ENTITY` and
    `WANDB_PROJECT` in the released `config.yaml`. No credential was ever
    exposed — those live in the `_wandb` blob and `wandb-metadata.json`, both
    already removed — but the published surface advertised a W&B account and
    project that are nothing to do with the recipe, and the sibling repo
    dropped a different set again. Both now drop the same one.
    A PPO `config.yaml` happens to hold its environment keys at the top level,
    so a top-level filter sufficed here today. That was luck rather than
    correctness — the same luck ran out one function down, where a nested
    `config_snapshot` published what a top-level filter could not see — so this
    recurses too, and stops depending on the shape of the file it is handed.
    """
    raw = yaml.safe_load(config_yaml.read_text())
    kept = drop_environment_keys(raw)
    config_yaml.write_text(yaml.safe_dump(kept, sort_keys=True))


def scrub_abs_paths(resume_json: Path) -> None:
    """Shorten absolute cluster paths and drop provenance from the sidecar.
    Shortening the snapshot alone left `wandb_run_id`, which
    save_checkpoint_metadata writes at the top level, in every released
    checkpoint's metadata. The sibling repo shipped the same id inside its
    pickled `.pth` files; both now drop it.
    Dropping it at the top level then left everything below it. `shorten_paths`
    recursed into `config_snapshot` while the key filter read only the top of
    the document, so `USE_WANDB`, `WANDB_PROJECT`, `WANDB_ENTITY` and
    `WANDB_DOWNLOAD_DIR` went out in both released checkpoints while this
    function reported a clean scrub. Both passes now recurse, by composition
    rather than by a special case for the one nesting we happened to know
    about, and the removed paths are named rather than assumed.
    """
    meta = json.loads(resume_json.read_text())
    dropped = environment_key_paths(meta)
    resume_json.write_text(json.dumps(scrub(meta), indent=2))
    if dropped:
        print(f"  scrubbed {', '.join(sorted(dropped))} from {resume_json.name}")


# =============================================================================
# Description
# =============================================================================

def describe(model_dir: Path, steps: list[int]) -> dict[str, str]:
    """Pull env name and training detail out of a checkpoint's own metadata."""
    resume = model_dir / "resume_metadata.json"
    if resume.exists():
        meta = json.loads(resume.read_text())
        cfg = meta["config_snapshot"]
        detail = f"{meta['total_gradient_steps_completed']:,} grad steps"
        arch = (
            f"{cfg['N_LAYERS']}L, d_model {cfg['D_MODEL']}, "
            f"{cfg['N_HEADS']} heads, horizon {cfg['PLAN_HORIZON']}"
        )
    else:
        raw = yaml.safe_load((model_dir / "config.yaml").read_text())
        cfg = {
            k: v["value"] for k, v in raw.items() if not is_environment_key(k)
        }
        detail = f"{float(cfg['TOTAL_TIMESTEPS']):.0e} frames"
        arch = f"RNN, layer size {cfg['LAYER_SIZE']}"
    return {
        "path": str(model_dir.relative_to(ROOT)),
        "role": ROLES.get(model_dir.parent.name, model_dir.parent.name),
        "env": cfg["ENV_NAME"],
        "arch": arch,
        "step": f"{max(steps):,}",
        "detail": detail,
        "size": human_size(model_dir),
    }


def describe_run(run: Path, staged: Path) -> dict[str, str]:
    """Summarise what an ablation run contributes to the release."""
    counts = [
        f"{len(list((staged / d).glob('*')))} {d}"
        for d in RUN_DIRS if (staged / d).is_dir()
    ]
    files = [f for f in RUN_FILES if (staged / f).is_file()]
    return {
        "run": run.name,
        "path": str(run.relative_to(ROOT)),
        "contents": ", ".join([*(f"`{f}`" for f in files), *counts]),
        "size": human_size(staged),
    }


def describe_inference(name: str, payload: dict) -> dict[str, str]:
    """Summarise one ``--mode inference`` result JSON."""
    metrics = payload.get("metrics", payload)
    score = metrics.get("mean_score")
    envs, steps = metrics.get("eval_num_envs"), metrics.get("eval_steps")
    return {
        "file": name,
        "env": payload.get("env_name", "-"),
        "episodes": f"{envs} envs x {steps} steps" if envs and steps else "-",
        "metric": (
            f"mean score {score:.2f}" if isinstance(score, int | float) else "-"
        ),
    }


# =============================================================================
# Staging
# =============================================================================

def stage_checkpoints(staging: Path, models: dict[Path, list[int]]) -> list[dict[str, str]]:
    """Copy each checkpoint directory, scrubbing its provenance metadata."""
    rows = []
    for model_dir, steps in models.items():
        target = staging / model_dir.relative_to(ROOT)
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copytree(model_dir, target, ignore=COPY_IGNORE)
        if (target / "config.yaml").exists():
            strip_wandb_block(target / "config.yaml")
        if (target / "resume_metadata.json").exists():
            scrub_abs_paths(target / "resume_metadata.json")
        scrub_staged_tree(target)
        rows.append(describe(model_dir, steps))
    return rows


def stage_runs(staging: Path, runs: list[Path]) -> list[dict[str, str]]:
    """Copy each ablation run's summary, tables and figures, scrubbed.
    These were copied verbatim, and every run's `results.json` embeds the
    config it ran under -- so `WANDB_ENTITY` went out in the published ablation
    suite. The checkpoint sidecar was scrubbed by name while this whole tree
    was not, which is the same defect one directory across.
    """
    rows = []
    for run in runs:
        target = staging / run.relative_to(ROOT)
        target.mkdir(parents=True, exist_ok=True)
        for name in RUN_FILES:
            if (run / name).is_file():
                shutil.copy2(run / name, target / name)
        for name in RUN_DIRS:
            if (run / name).is_dir():
                shutil.copytree(run / name, target / name, ignore=COPY_IGNORE)
        scrub_staged_tree(target)
        rows.append(describe_run(run, target))
    return rows


def stage_inference(staging: Path, files: list[Path]) -> list[dict[str, str]]:
    """Copy each inference result JSON into ``results/inference/``.
    Scrubbed, not merely shortened: this staged paths without ever filtering
    environment keys, which is the same asymmetry that leaked from the
    checkpoint sidecar, one step further along. No inference payload carries an
    environment key today; the point is that nothing here depends on that.
    """
    target_dir = staging / INFERENCE.relative_to(ROOT)
    rows: list[dict[str, str]] = []
    for src in files:
        try:
            payload = json.loads(src.read_text())
        except json.JSONDecodeError:
            print(f"Skipping unreadable inference result {src}.", file=sys.stderr)
            continue
        name = src.name
        if any(r["file"] == name for r in rows):
            name = f"{src.parent.name}-{src.name}"
        target_dir.mkdir(parents=True, exist_ok=True)
        (target_dir / name).write_text(
            json.dumps(scrub(payload), indent=2) + "\n",
        )
        row = describe_inference(name, payload)
        row["size"] = human_size(target_dir / name)
        rows.append(row)
    return rows


def stage_paper_figures(staging: Path, figures: list[Path]) -> list[dict[str, str]]:
    """Copy the manuscript figures into ``results/paper_figures/``."""
    if not figures:
        return []
    target_dir = staging / PAPER_FIGURES.relative_to(ROOT)
    target_dir.mkdir(parents=True, exist_ok=True)
    rows = []
    for src in figures:
        shutil.copy2(src, target_dir / src.name)
        rows.append({"file": src.name, "size": human_size(target_dir / src.name)})
    return rows


def stage(
    staging: Path,
    models: dict[Path, list[int]],
    runs: list[Path],
    inference: list[Path],
    paper_figures: list[Path],
) -> tuple[
    list[dict[str, str]],
    list[dict[str, str]],
    list[dict[str, str]],
    list[dict[str, str]],
]:
    """Stage checkpoints, results and LICENSE; the card is written by the caller.
    LICENSE comes from git rather than the working tree: a
    ``hf download --local-dir .`` overwrites it with the Hub's own copy, and
    publishing from the tree would push that straight back up as current.
    """
    rows = stage_checkpoints(staging, models)
    run_rows = stage_runs(staging, runs)
    inf_rows = stage_inference(staging, inference)
    fig_rows = stage_paper_figures(staging, paper_figures)
    # ROOT is passed explicitly rather than read inside the helper, so a
    # test can rebind it and have that take effect.
    copy_tracked_file("LICENSE", staging / "LICENSE", ROOT)
    return rows, run_rows, inf_rows, fig_rows


# =============================================================================
# Model card
# =============================================================================

def table(headers: list[str], lines: list[str]) -> str:
    sep = "|".join(["---"] * len(headers))
    return f"| {' | '.join(headers)} |\n|{sep}|\n" + "".join(f"| {ln} |\n" for ln in lines)


def checkpoint_table(rows: list[dict[str, str]]) -> str:
    return table(
        ["Path", "Role", "Environment", "Architecture", "Selected at", "Training", "Size"],
        [
            f"`{r['path']}` | {r['role']} | `{r['env']}` | {r['arch']} | "
            f"{r['step']} | {r['detail']} | {r['size']}"
            for r in sorted(rows, key=lambda r: r["path"])
        ],
    )


def results_section(
    run_rows: list[dict[str, str]],
    inf_rows: list[dict[str, str]],
    fig_rows: list[dict[str, str]],
) -> str:
    """Ablation, inference and manuscript-figure tables; empty when none."""
    parts = []
    if run_rows:
        parts.append(
            "RL fine-tuning ablation runs, as produced by "
            "`experiments/rl_finetuning/run_ablations.py`. Each run ships its "
            "`results.json` summary, the `diagnosis.md` write-up, and the "
            "tables (`.csv` and `.tex`) and figures generated from it.\n\n"
            + table(
                ["Run", "Contents", "Size"],
                [
                    f"`{r['path']}` | {r['contents']} | {r['size']}"
                    for r in sorted(run_rows, key=lambda r: r["run"])
                ],
            ),
        )
    if inf_rows:
        parts.append(
            "Evaluation results produced by `main.py --mode inference` on the "
            "checkpoints above, under `results/inference/`.\n\n"
            + table(
                ["File", "Environment", "Evaluation", "Headline metric", "Size"],
                [
                    f"`{r['file']}` | `{r['env']}` | {r['episodes']} | "
                    f"{r['metric']} | {r['size']}"
                    for r in sorted(inf_rows, key=lambda r: r["file"])
                ],
            ),
        )
    if fig_rows:
        parts.append(
            "Manuscript figures, as vector PDF at NeurIPS column width, under "
            "`results/paper_figures/`. These are built by "
            "`scripts/paper_figures.py`, which reads the ablation "
            "`results.json` of *both* environments and draws Craftax Classic "
            "and MiniHack side by side, so the identical set is published in "
            "this release and in the MiniHack one.\n\n"
            + table(
                ["Figure", "Size"],
                [
                    f"`{r['file']}` | {r['size']}"
                    for r in sorted(fig_rows, key=lambda r: r["file"])
                ],
            ),
        )
    return "## Results\n\n" + "\n".join(parts) if parts else ""


def featured(rows: list[dict[str, str]]) -> dict[str, str]:
    """The checkpoint the download and usage examples are written against."""
    planners = [r for r in rows if "planner" in r["role"].lower()]
    return sorted(planners or rows, key=lambda r: r["path"])[0]


def model_card(
    repo_id: str,
    rows: list[dict[str, str]],
    run_rows: list[dict[str, str]],
    inf_rows: list[dict[str, str]],
    fig_rows: list[dict[str, str]],
    total_mb: float,
) -> str:
    example = featured(rows)
    return f"""---
license: mit
library_name: jax
pipeline_tag: reinforcement-learning
tags:
- reinforcement-learning
- planning
- discrete-diffusion
- remdm
- craftax
- jax
- flax
- orbax
---
# ReMDM Planner: {ENV_NAME} checkpoints
Trained weights accompanying *{PAPER}*: a remasking discrete diffusion model
(ReMDM) used as an action-sequence planner in
[Craftax](https://github.com/MichaelTMatthews/Craftax), together with the
PPO-RNN experts that supervise it, and the results reported in the paper.
Code, configs and evaluation harness: {CODE_URL}
## Contents
{checkpoint_table(rows)}
Each diffusion checkpoint ships a `resume_metadata.json` holding the full
config snapshot it was trained under; each PPO expert ships `config.yaml` and
`wandb-summary.json` (final training metrics).
Weights are [Orbax](https://orbax.readthedocs.io) checkpoint directories
(OCDBT format), not `safetensors` — the models are Flax modules restored via
`orbax.checkpoint`, and the paths above mirror the source repository so a
snapshot can be dropped straight into a working copy.
{results_section(run_rows, inf_rows, fig_rows)}
## Download
This repo mirrors the code repository's layout, so a snapshot drops straight
into a working copy -- but it also carries its own `README.md` (this card),
`LICENSE` and `.gitattributes`, and `local_dir="."` would overwrite the code
repository's copies of all three. Exclude them, or download into a directory
of its own.
```python
from huggingface_hub import snapshot_download
# everything (~{total_mb:.0f} MB), into a clone of the code repository
snapshot_download(
    repo_id="{repo_id}",
    local_dir=".",
    ignore_patterns=["README.md", "LICENSE", ".gitattributes"],
)
# or somewhere of its own, leaving any working copy untouched
snapshot_download(repo_id="{repo_id}", local_dir="remdm-planner-craftax")
# a single model
snapshot_download(
    repo_id="{repo_id}",
    local_dir=".",
    allow_patterns="{example['path']}/**",
)
```
## Use
From a clone of the code repository, after downloading into it:
```bash
uv run python main.py --mode inference \\
    --checkpoint {example['path']} \\
    --output results/inference/eval.json
```
Programmatic loading uses `src.planners.model.load_checkpoint` for the
diffusion planners and `src.planners.ppo.load_ppo_agent` for the experts; both
take the checkpoint directory path and restore the latest step. Architecture
arguments should be read from the checkpoint's own `resume_metadata.json`
rather than hardcoded.
## Training
The diffusion planners are bidirectional transformers that denoise a masked
action plan conditioned on the symbolic observation, trained either by offline
behaviour cloning on PPO rollouts or by online DAgger against the PPO expert.
Model size and horizon differ per run (see the table); the PPO-RNN experts are
the Craftax baselines. Exact hyperparameters for every run, including the
remasking strategy, schedule and sampling settings, are in the per-checkpoint
metadata files listed above, which are the authoritative record.
Directory names encode the environment and the total environment timesteps the
run was trained for. `Selected at` is whatever each run used as its Orbax step
counter, which is environment frames for the runs published here.
## Limitations
These are research artefacts tied to specific Craftax versions and symbolic
observation encodings; they are not general-purpose agents and will not
transfer to other environments or to pixel observations. Evaluation results and
their variance are reported in the paper.
## Citation
```bibtex
@inproceedings{{remdm-planner-craftax-planner,
  title  = {{{PAPER}}},
  author = {{Anonymous}},
  year   = {{2026}},
  note   = {{NeurIPS 2026 Workshop: Beyond Next-Token Prediction}}
}}
```
## License
MIT, see `LICENSE`.
"""


# =============================================================================
# Entry point
# =============================================================================

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    p.add_argument("--repo-id", required=True, help="e.g. <ANON_HF_REPO_ID>")
    p.add_argument(
        "--inference-results", nargs="*", default=[], metavar="PATH",
        help=f"extra --mode inference JSONs or directories, on top of "
             f"{INFERENCE.relative_to(ROOT)}/",
    )
    p.add_argument("--private", action="store_true", help="create the repo private")
    p.add_argument("--dry-run", action="store_true", help="stage and print, do not upload")
    p.add_argument("--yes", action="store_true", help="skip the confirmation prompt")
    return p.parse_args()


def main() -> int:
    args = parse_args()

    token = os.environ.get("HF_TOKEN")
    if not args.dry_run and not token:
        print("HF_TOKEN is not set.", file=sys.stderr)
        return 1

    models = discover_checkpoints()
    if not models:
        print(
            f"No Orbax checkpoints found under {CKPTS}.\n"
            "Discovery expects the released layout, "
            "checkpoints/<role>/<name>/<step>/, and skips checkpoints/hf/ "
            "because that is where Hub downloads land. Copy a run's "
            "`wandb.run.dir/policies` directory to "
            "checkpoints/{offline,online}/<name> first.",
            file=sys.stderr,
        )
        return 1
    runs = discover_runs()
    inference = discover_inference(args.inference_results)
    paper_figures = discover_paper_figures()

    with tempfile.TemporaryDirectory(prefix="remdm-planner-craftax-") as tmp:
        staging = Path(tmp)
        rows, run_rows, inf_rows, fig_rows = stage(
            staging, models, runs, inference, paper_figures
        )
        total_mb = dir_size_mb(staging)
        card = model_card(
            args.repo_id, rows, run_rows, inf_rows, fig_rows, total_mb
        )
        (staging / "README.md").write_text(card)

        files = [f for f in staging.rglob("*") if f.is_file()]
        print(f"Staged {plural(len(rows), 'checkpoint')}, "
              f"{plural(len(run_rows), 'ablation run')}, "
              f"{plural(len(inf_rows), 'inference result')}, "
              f"{plural(len(fig_rows), 'paper figure')}, "
              f"{plural(len(files), 'file')}, {total_mb:.0f} MB")
        for r in sorted(rows, key=lambda r: r["path"]):
            print(f"  {r['path']:<70} {r['size']:>8}")
        for r in sorted(run_rows, key=lambda r: r["run"]):
            print(f"  {r['path']:<70} {r['size']:>8}  {r['contents']}")
        for r in sorted(inf_rows, key=lambda r: r["file"]):
            print(f"  results/inference/{r['file']:<52} {r['size']:>8}  {r['metric']}")
        for r in sorted(fig_rows, key=lambda r: r["file"]):
            print(f"  results/paper_figures/{r['file']:<48} {r['size']:>8}")
        if not run_rows:
            print(f"Warning: no ablation runs with a results.json under {RUNS}.",
                  file=sys.stderr)
        if not inf_rows:
            print("Warning: no inference results; produce them with "
                  "`main.py --mode inference --output "
                  f"{INFERENCE.relative_to(ROOT)}/<name>.json`.", file=sys.stderr)

        if not fig_rows:
            print("Warning: no manuscript figures; build them with "
                  "`uv run python scripts/paper_figures.py --outdir "
                  f"{PAPER_FIGURES.relative_to(ROOT)}` once the MiniHack "
                  "sibling's results.json is present.", file=sys.stderr)

        if args.dry_run:
            print(f"Dry run; staged tree left nowhere. Card:\n\n{card}")
            return 0

        if not args.yes:
            visibility = "private" if args.private else "public"
            if input(f"Upload to {args.repo_id} ({visibility})? [y/N] ").strip().lower() not in {"y", "yes"}:
                print("Aborted.")
                return 0

        from huggingface_hub import HfApi

        api = HfApi(token=token)
        api.create_repo(args.repo_id, repo_type="model", private=args.private, exist_ok=True)
        api.upload_folder(
            repo_id=args.repo_id,
            folder_path=str(staging),
            repo_type="model",
            ignore_patterns=HUB_IGNORE,
            commit_message="Upload Craftax ReMDM planner checkpoints and results",
        )
        print(f"Done: https://huggingface.co/{args.repo_id}")
    return 0


if __name__ == "__main__":
    sys.exit(main())