[RyanPCo]

Slurm requeue detection moves above ModelWrapper construction and reads the
checkpoint dir from cfg (callbacks.model_checkpoint.dirpath, else
/checkpoints); it resumes only when last.ckpt exists (a job preempted
before its first checkpoint warns and falls back to the launch-time ckpt_path).
When cfg.ckpt_path is an existing checkpoint file, the model's config tree gets
model.robomimic_model.config.pytorch_weight_path=null: the checkpoint holds
every weight, so a PI run resumes or evaluates without the multi-GB base
safetensors on disk. cfg itself keeps the path, so logged hyperparameters are
unchanged. Lightning's special values (last, best, hpc, registry) keep the base
weights, since they may resolve to no checkpoint. Exception: eval_latent's
pretrained=true, where ckpt_path only routes the output dir and the base
weights are the model under test.

Eval loads (trainHydra eval branch, eval_latent.run) go through
load_checkpoint_weights (mmap'd torch.load): missing keys raise, unexpected
keys warn. With base weights skipped, strict=False would have run random
weights silently.

Co-Authored-By: Claude Fable 5.1 [email protected]
Claude-Session: https://claude.ai/code/session_01MWtR1x5tjREwoxQZAnUGrh
Co-Authored-By: Claude Opus 5 (1M context) [email protected]
Claude-Session: https://claude.ai/code/session_01BbQNNaMAbTNMypegmJkv4s

🤖 Generated with Claude Code

https://claude.ai/code/session_01MWtR1x5tjREwoxQZAnUGrh
https://claude.ai/code/session_01BbQNNaMAbTNMypegmJkv4s

[RyanPCo]

• train: give each Slurm slurm job gets its own cache dir to avoid conflict #638 : 2 dependent PRs (#639 , #653 )
• write git version into the run dir for reproducibility #637
• train: don't load PI base weights on resume and verify that checkpoint path exists #636 👈 (View in Graphite)
• configs: declare proprio and actions dims once per HPT model config #634 : 1 other dependent PR (#635 )
• cache norm stats for data configs to minimize recomputation #633
• import cloud packages lazily to support local training #631 : 1 other dependent PR (#632 )
• cache dataset resolver for faster training initialization #630
• validate episode hashes in data config before training #628
• modify pi camera map to match data #629
• add tests to ensure training works between every model and data vendor #626
• run CI tests on every PR #625
• add env variables for user-specific paths #624
• build: set pyproject and uv as the dependency managers #623
• fix(configs): fix paths for data hydra yaml configs #622
• refactor(algo): delete ACT #627 : 1 other dependent PR (#621 )
• tests: create dedicated test folder with unit and integration tests #620
• add pyright type checker #619
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[RyanPCo]

@greptileai

[greptile-apps]

Confidence Score: 4/5

[Medium risk] Refactors checkpoint loading and model configuration initialization.

The PR should not merge until evaluation can load checkpoints that the existing training loader supports.

Findings

1.  Legacy checkpoints fail evaluation ▶
2.  Shared width ignores action padding ▶

Summary

The PR selects an existing checkpoint before constructing the model, skips PI base-weight loading when checkpoint weights will replace it, and checks for missing keys during evaluation. It also centralizes HPT widths in dims and resolves interpolations when logging hyperparameters.

• Evaluation needs the existing fallback for checkpoints that cannot be memory-mapped.
• Shared-head width overrides need to agree with action preprocessing.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Composed config] --> B{Slurm requeue with last.ckpt?}
  B -- Yes --> C[Select last.ckpt]
  B -- No --> D[Keep launch-time ckpt_path]
  C --> E{Existing checkpoint file and not pretrained?}
  D --> E
  E -- Yes --> F[Construct PI without base weights]
  E -- No --> G[Construct PI with configured base weights]
  F --> H[Restore checkpoint for train or eval]
  G --> H

Loading

_{Reviews (1) · Last reviewed commit: "train: settle ckpt_path before the model..."}

[greptile-apps]

Legacy checkpoints fail evaluation

If an evaluation checkpoint uses PyTorch's legacy, non-zip format, mmap=True raises before any weights are loaded. The training checkpoint loader falls back to a normal read for this format, but this new evaluation loader does not. As a result, a checkpoint that can resume training cannot be evaluated. Please provide the same fallback here.

[greptile-apps]

Shared width ignores action padding

This config, along with the other shared-head configs, now uses dims.action_width for sampling and the model input, but action preprocessing still pads 12-dimensional human actions to exactly 14 dimensions. Overriding action_width to 16 would therefore send 14-dimensional training inputs to a 16-dimensional head. Either make padding follow the configured width or document and constrain the width to 14.

[RyanPCo]

Merge activity

• Oct 19, 7:23 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• Oct 19, 7:25 PM UTC: Graphite rebased this pull request as part of a merge.
• Oct 19, 7:25 PM UTC: @RyanPCo merged this pull request with Graphite.