"""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 \\ [--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////`` — 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"(? 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. ") 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////, and skips checkpoints/hf/ " "because that is where Hub downloads land. Copy a run's " "`wandb.run.dir/policies` directory to " "checkpoints/{offline,online}/ 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)}/.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())