How to evaluate Hugging Face per‑task checkpoints (e.g. `walker-walk.pt`) with this repo?

[keyboardAnt]

Hi @nicklashansen, thanks a lot for releasing Newt, MMBench, and the checkpoints.

I’m trying to run a sanity‑check evaluation of one of the per‑task checkpoints from the HF repo (https://huggingface.co/nicklashansen/newt/blob/main/walker-walk.pt) using the current code in this repo (commit d0cea95).

I ran:

cd tdmpc2

python train.py \
  task=walker-walk \
  model_size=B \
  checkpoint=/path/to/hf_checkpoints/walker-walk.pt \
  steps=1 \
  num_envs=2 \
  use_demos=False \
  tasks_fp=/path/to/newt/tasks.json \
  exp_name=eval_hf_walker_walk \
  save_video=True \
  env_mode=sync \
  compile=False

This fails with:

AssertionError: pad should be positive
...
File "tdmpc2/common/layers.py", line 190, in api_model_conversion
    assert pad > 0, 'pad should be positive'

I inspected the HF checkpoint:

state = torch.load("hf_checkpoints/walker-walk.pt", map_location="cpu", weights_only=False)
state = state["model"]

print("_task_emb.weight", state["_task_emb.weight"].shape)       # torch.Size([10, 512])
print("_action_masks", state["_action_masks"].shape)             # torch.Size([10, 7])
print("_encoder.state.0.weight", state["_encoder.state.0.weight"].shape)  # torch.Size([256, 554])
print("_dynamics.0.weight", state["_dynamics.0.weight"].shape)   # torch.Size([512, 1031])

From this I infer the HF walker-walk.pt was trained with something like:

• model_size = B (latent_dim=512, enc_dim=256, mlp_dim=512),
• task_dim = 512 (10 tasks × 512‑dim embedding),
• obs_state_dim = 42 (since 554 = 42 + 512),
• action_dim = 7 (since 1031 = 512 (latent) + 7 (action) + 512 (task)).

However, the current single‑task path in this repo builds a model with:

• task != "soup" → task_dim = 0 in parse_cfg,
• padded state obs_shape['state'] = (128,) via VecWrapper,
• padded actions (16,) via VecWrapper.

So the local _encoder.state.0.weight is [256, 128], which is smaller than the HF [256, 554], and api_model_conversion’s assumption (“target has more input channels than source, pad source”) is violated, hence pad < 0 and the assertion.

Environment: Docker image built from this repo’s docker/Dockerfile.

---

Questions

1. What is the exact config used to train the per‑task HF checkpoints, e.g. walker-walk.pt?

   • model_size?
   • obs (state vs rgb vs state+rgb)?
   • task_dim?
   • How many tasks / which tasks correspond to the 10 rows in _task_emb.weight and _action_masks?

2. What is the recommended way to evaluate these HF per‑task checkpoints with this implementation?

   • Is there a matching config/script you use internally (e.g. a specific train.py invocation or a separate eval script)?
   • Should we be using a multitask (task="soup" or similar) config with task_dim=512 and then fix the task index at eval time, rather than the current task_dim=0 single‑task path?

3. More generally: are the HF single‑task .pt files intended to be evaluated with this early code release as‑is, or should we wait for a dedicated evaluation script / config that matches those checkpoints?

Any guidance or an example eval command for walker-walk.pt (or any other per‑task checkpoint) would be greatly appreciated.

[nicklashansen]

Thanks for raising this issue!

are the HF single‑task .pt files intended to be evaluated with this early code release as‑is, or should we wait for a dedicated evaluation script / config that matches those checkpoints?

I am planning to write a dedicated evaluation script to streamline things + readme instructions for how to run evals, your issue has prompted me to get this done asap, I'll get back to you soon!

In the meantime, you should be able to repurpose generate_demos.py for evaluation of single-task checkpoints. Example usage of this script would be: python generate_demos.py task=walker-walk after you set CHECKPOINT_PATH in the generate_demos.py file to your checkpoint directory.

Apologies that some aspects of the codebase are still a bit of a work in progress!

[nicklashansen]

Ok I had a chance to look into this and I was able to reproduce your error! It looks like I might have uploaded some of the checkpoints a bit hastily; this checkpoint originated from an earlier (internal) set of single-task runs that indeed have other dims. I'll need a bit of time to iterate through all 200 checkpoints and replace the incompatible ones with the newer, correct checkpoints but I'm on it :-) I checked the multitask checkpoints and they all load correctly on my end!

[nicklashansen]

@keyboardAnt I have reuploaded all single-task checkpoints and am able to load them on my end, so you should be able to do so as well! All the reuploaded checkpoints now have 128 dim obs, 16 dim actions, and 72.3MB file size. I have also updated config.py to automatically set model size etc. when in single-task mode, this should make the code a little bit easier to figure out. You can now load and evaluate a checkpoint as follows:

python train.py wandb_project=<project> wandb_entity=<entity> checkpoint="<path>/walker-walk.pt" task=walker-walk

A dedicated evaluation script (as well as huggingface model api potentially) is still on my to-do list. Thanks for your patience! Let me know if you run into any other issues :-)

[keyboardAnt]

Thanks @nicklashansen. I've been using this script for eval-only. Is it similar to what you had in mind?

[nicklashansen]

Yes that looks like a very reasonable workflow! One thing to watch out for depending on how robust you want your evals to be: num_envs and eval_episodes control the number of evaluations you are doing (actual number of episodes will be num_envs * eval_episodes). Some tasks can have quite a bit of variance between episodes so you can consider upping those numbers a bit