Enable MolmoAct2 rollout on SO-100/101 with calibration correction (#3879)

* fix(rollout): improve visual feature mismatch error with --rename_map hint

* feat(policies): add joint frame transform and hardware deployment docs for MolmoAct2

Add MolmoAct2StateFrameTransformStep and MolmoAct2ActionFrameTransformStep
processor steps for cross-calibration compatibility on SO-100/101. Add
joint_signs and joint_offsets config fields. Add hardware deployment section
to molmoact2.mdx with camera naming convention, joint frame correction, and
safety guidance.

* chore(docs): address PR comment

* fix: address reviewer comments

AGENT_GUIDE.md

@@ -138,7 +138,7 @@ lerobot-replay --robot.type=so101_follower --robot.port=<FOLLOWER_PORT> --robot.
   --dataset.repo_id=${HF_USER}/my_task --dataset.episode=0
 ```
 
-**4.9 Train** (default: ACT — fastest, lowest memory). Apple silicon: `--policy.device=mps`. See §6/§7 for policy and duration.
+**4.9 Train** (default: ACT — fastest, lowest memory). Apple silicon: `--policy.device=mps`. No local GPU? Add `--job.target=<flavor>` (e.g. `a10g-small`, list them with `hf jobs hardware`) to run on Hugging Face Jobs instead. See §6/§7 for policy and duration.
 
 ```bash
 lerobot-train \

docs/source/cheat-sheet.mdx

@@ -150,6 +150,14 @@ lerobot-train \
     --steps=20000
 ```
 
+No local GPU? Add `--job.target=<flavor>` (e.g. `a10g-small`) to either command and `lerobot-train` runs it on [Hugging Face Jobs](https://huggingface.co/docs/hub/jobs) instead — it uploads a local-only dataset for you and pushes the trained model. List flavors with `hf jobs hardware`.
+
+To resume, point `--config_path` at a checkpoint and add `--resume=true`. It accepts a local path or a Hub repo id (the latest checkpoint is fetched), and works locally or on a job by adding `--job.target=<flavor>`:
+
+```bash
+lerobot-train --config_path=${HF_USER}/policy_test --resume=true --job.target=a10g-small
+```
+
 ### Inference
 
 Inference means running the trained policy/model on a robot. For that we use `lerobot-rollout`. You will need to provide a path to your policy. It can be a local path or a path to Hugging Face for example "lerobot/folding_latest". Your cameras configuration needs to match what was used when collecting the dataset. Duration is in seconds if unspecified, it will run forever.

docs/source/hardware_guide.mdx

@@ -96,3 +96,4 @@ Notes:
 - The leading `nvidia-smi` is a quick sanity check that CUDA is visible inside the container — useful to fail fast if the flavor or driver mismatched.
 - The default Job timeout is 30 minutes; pass `--timeout 4h` (or longer) for real training.
 - `--flavor` maps onto the table above: `t4-small`/`t4-medium` (T4, ACT only), `l4x1`/`l4x4` (L4 24 GB), `a10g-small/large/largex2/largex4` (A10G 24 GB scaled out), `a100-large` (A100). For the current full catalogue + pricing see [https://huggingface.co/docs/hub/jobs](https://huggingface.co/docs/hub/jobs).
+- Prefer not to write the `hf jobs run` wrapper yourself? `lerobot-train` can submit the job for you: just add `--job.target=<flavor>` to a normal training command and it handles dataset upload, log streaming, and the final model push. See the [imitation-learning training guide](./il_robots).

docs/source/il_robots.mdx

@@ -514,6 +514,12 @@ lerobot-train \
   --resume=true
 ```
 
+`--config_path` also accepts a **Hub repo id**: if a run pushed its checkpoints to the Hub (with `--save_checkpoint_to_hub=true`), you can resume straight from the repo — its latest checkpoint is downloaded and training continues, restoring the optimizer, scheduler, step counter and data order:
+
+```bash
+lerobot-train --config_path=${HF_USER}/my_policy --resume=true
+```
+
 If you do not want to push your model to the hub after training use `--policy.push_to_hub=false`.
 
 Additionally you can provide extra `tags` or specify a `license` for your model or make the model repo `private` by adding this: `--policy.private=true --policy.tags=\[ppo,rl\] --policy.license=mit`
@@ -526,7 +532,9 @@ If your local computer doesn't have a powerful GPU you could utilize Google Cola
 
 Hugging Face jobs let's you easily select hardware and run the training in the cloud. So if you don't have a powerful GPU or you need more VRAM or just want to train a model much faster use HF Jobs! It's pay as you go and you simply pay for each second of use, you can see the pricing and additional information [here](https://huggingface.co/docs/hub/jobs).
 
-To run the training use this command:
+> **Tip:** if you just want to launch a standard training run, you can skip building the command below and use the integrated **Train on HF Jobs via `--job.target`** flow described further down — `lerobot-train` then submits the job, uploads a local-only dataset for you, and streams the logs.
+
+To run the training manually use this command:
 
 <hfoptions id="train_with_hf_jobs">
 <hfoption id="Command">
@@ -599,6 +607,51 @@ Once the training is started you can go to [Jobs](https://huggingface.co/setting
 
 After training the model will be pushed to hub and you can use it as any other model with LeRobot.
 
+#### Train on HF Jobs via `--job.target` (integrated CLI)
+
+`lerobot-train` runs locally by default. To run on a HuggingFace GPU without constructing the Docker command yourself, pass `--job.target` with a hardware flavor name:
+
+```bash
+lerobot-train \
+  --dataset.repo_id=${HF_USER}/so101_test \
+  --policy.type=act \
+  --policy.repo_id=${HF_USER}/my_policy \
+  --job.target=a10g-small
+```
+
+List available flavors and pricing with `hf jobs hardware`. The run streams its logs to your terminal; press Ctrl-C to detach (the job keeps running in the cloud). Re-attach or cancel with:
+
+```bash
+hf jobs logs <job-id>
+hf jobs cancel <job-id>
+```
+
+If your dataset exists only locally (not yet on the Hub), it is automatically pushed to a **private** Hub repo so the job can download it by `repo_id` (nothing is made public). The trained model is pushed to the model repo at the end of the run. To also push every intermediate checkpoint to the Hub as it is saved (so you can monitor progress mid-run), add `--save_checkpoint_to_hub=true` — this requires a runtime image that includes this feature.
+
+Every job (and any dataset pushed by the run) is tagged `lerobot` so it's easy to find on the Hub. Add your own with `--job.tags '["my-tag"]'`.
+
+By default the job is capped at `2d` (48h) of wall-clock. Override it with an HF Jobs duration string, e.g. `--job.timeout=4h` to fail faster or `--job.timeout=7d` for a longer run.
+
+> **Note:** the model repo is created up front (it holds the staged training config the job runs from). If a run fails before the model is pushed, that repo is left on the Hub so you can inspect it — it is not deleted automatically, so repeated failures can leave empty repos behind. Remove one with `hf repo delete <repo-id>`.
+
+**Prerequisites:** run `hf auth login` before submitting. For Weights & Biases integration, run `wandb login` or set `WANDB_API_KEY` on your machine — the key is forwarded to the job automatically.
+
+**Resuming on a job.** Adding `--job.target` to a resume command runs the resume in the cloud — the same command works locally or remotely. The checkpoint repo is the source of truth, and new checkpoints continue the lineage in the same repo:
+
+```bash
+# resume a Hub run on a job (its checkpoints are already on the Hub)
+lerobot-train --config_path=${HF_USER}/my_policy --resume=true --job.target=a10g-small
+
+# resume a LOCAL run on a job — the checkpoint is uploaded to a private Hub repo first,
+# then the job resumes from it (a local-only dataset is uploaded the same way)
+lerobot-train \
+  --config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \
+  --resume=true \
+  --job.target=a10g-small
+```
+
+Job settings come from the current command, so override `--job.target`, `--job.timeout`, etc. as needed; for the resumed run to itself be resumable later, keep `--save_checkpoint_to_hub=true`.
+
 #### Upload policy checkpoints
 
 Once training is done, upload the latest checkpoint with:
@@ -620,6 +673,8 @@ hf upload ${HF_USER}/act_so101_test${CKPT} \
 
 Use `lerobot-rollout` to deploy a trained policy on your robot. You can choose different strategies depending on your needs:
 
+The examples below load the model from `--policy.path`. To pin a specific pushed version — useful once `--save_checkpoint_to_hub=true` has committed several checkpoints — add `--policy.pretrained_revision` with a commit hash, branch, or tag. Each pushed checkpoint is tagged with its step (e.g. `--policy.pretrained_revision=010000`), so you can recover a checkpoint by step without looking up its commit sha.
+
 <hfoptions id="eval">
 <hfoption id="Base mode (no recording)">
 ```bash

docs/source/molmoact2.mdx

@@ -386,6 +386,68 @@ These results demonstrate MolmoAct2's strong performance across diverse robotic
 manipulation tasks. To reproduce them, follow the instructions in the LIBERO
 evaluation section.
 
+## Hardware Deployment (lerobot-rollout)
+
+LeRobot-format checkpoints are available on the Hub for direct use with
+`lerobot-rollout`. Each checkpoint uses specific camera names that must
+match your robot's camera configuration.
+
+### Camera naming convention
+
+Each checkpoint expects specific `observation.images.*` keys.
+If your robot cameras have different names, use `--rename_map` to map them:
+
+| Checkpoint                    | Camera keys            | Description              |
+| ----------------------------- | ---------------------- | ------------------------ |
+| MolmoAct2-LIBERO-LeRobot      | `image`, `wrist_image` | LIBERO sim cameras       |
+| MolmoAct2-BimanualYAM-LeRobot | `top`, `left`, `right` | YAM 3-camera setup       |
+| MolmoAct2-DROID-LeRobot       | `cam0`, `cam1`         | External + wrist         |
+| MolmoAct2-SO100_101-LeRobot   | `cam0`, `cam1`         | Primary + secondary view |
+
+Example with an SO-100 robot using top and side cameras:
+
+```bash
+lerobot-rollout \
+  --policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \
+  --rename_map='{"observation.images.top": "observation.images.cam0", "observation.images.side": "observation.images.cam1"}' \
+  --robot.type=so100_follower \
+  --robot.port=/dev/ttyACM0 \
+  --robot.cameras='{
+      top: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30},
+      side: {type: opencv, index_or_path: 2, width: 640, height: 480, fps: 30}
+  }' \
+  --task="pick up the red cube" --duration=30
+```
+
+To use a wrist camera instead, just change the rename mapping:
+
+```bash
+--rename_map='{"observation.images.top": "observation.images.cam0", "observation.images.wrist": "observation.images.cam1"}'
+```
+
+### Joint frame transform (SO-100/101 zero-shot)
+
+<Tip warning={true}>
+The MolmoAct2-SO100_101 checkpoint was trained on data that uses a different
+joint calibration convention than LeRobot >= 0.5.0. Without a frame
+correction, the arm may move in the wrong direction.
+
+This affects both **zero-shot deployment** and **fine-tuning** from the
+original checkpoint. The pretrained weights expect the old convention, so
+all joint data (observations and actions) must be transformed to match.
+
+The converted LeRobot checkpoint (`lerobot/MolmoAct2-SO100_101-LeRobot`)
+already includes this correction in its processor pipeline. If you convert
+or fine-tune the checkpoint yourself, set the following in the policy config (`configuration_molmoact2.py`):
+
+- `joint_signs`: `[1, -1, 1, 1, 1, 1]` (flips shoulder_lift direction)
+- `joint_offsets`: `[0, 90, 90, 0, 0, 0]` (shifts shoulder_lift and elbow_flex by 90°)
+
+See the [backward compatibility guide](./backwardcomp) for details on the
+calibration change.
+
+</Tip>
+
 ## Differences From the Original Implementation
 
 This LeRobot port is intended to match MolmoAct2 behavior while using LeRobot's

pyproject.toml

@@ -124,7 +124,7 @@ hardware = [
     "lerobot[deepdiff-dep]",
 ]
 viz = [
-    "rerun-sdk>=0.24.0,<0.27.0",
+    "rerun-sdk>=0.24.0,<0.34.0",
 ]
 # ── User-facing composite extras (map to CLI scripts) ─────
 # lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc.

src/lerobot/common/train_utils.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 from pathlib import Path
 
+from huggingface_hub import HfApi, snapshot_download
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
@@ -35,6 +36,7 @@ from lerobot.utils.constants import (
     TRAINING_STATE_DIR,
     TRAINING_STEP,
 )
+from lerobot.utils.hub import find_latest_hub_checkpoint
 from lerobot.utils.io_utils import load_json, write_json
 from lerobot.utils.random_utils import load_rng_state, save_rng_state
 
@@ -283,3 +285,61 @@ def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None:
     with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg):
         sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd)
     optimizer.load_state_dict(sharded_osd)
+
+
+def push_checkpoint_to_hub(
+    checkpoint_dir: Path,
+    repo_id: str,
+    *,
+    private: bool | None = None,
+) -> None:
+    """Upload a saved checkpoint directory to the Hub under checkpoints/<name>/.
+
+    Called once per save step when save_checkpoint_to_hub is enabled, so a
+    timed-out or crashed run still leaves recoverable checkpoints on the Hub.
+    The model repo is created idempotently, and the commit is tagged with the
+    checkpoint step so a checkpoint can be recovered with
+    --policy.pretrained_revision=<step> instead of a commit sha.
+    """
+    api = HfApi()
+    api.create_repo(repo_id=repo_id, repo_type="model", private=private, exist_ok=True)
+    commit = api.upload_folder(
+        folder_path=str(checkpoint_dir),
+        repo_id=repo_id,
+        repo_type="model",
+        path_in_repo=f"checkpoints/{checkpoint_dir.name}",
+        commit_message=f"checkpoint {checkpoint_dir.name}",
+    )
+    api.create_tag(
+        repo_id=repo_id,
+        tag=checkpoint_dir.name,
+        revision=commit.oid,
+        repo_type="model",
+        exist_ok=True,
+    )
+
+
+def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path:
+    """Download the latest checkpoint of a Hub training repo into a local run dir.
+
+    The symmetric counterpart to `push_checkpoint_to_hub`: given a model repo holding
+    `checkpoints/<step>/{pretrained_model,training_state}` subtrees, download the highest-numbered step
+    into `output_dir/checkpoints/<step>/`, recreate the local `last` symlink, and return that local
+    checkpoint dir. Used to resume training from the Hub on a machine (or HF Jobs pod) that does not
+    have the original local run dir.
+    """
+    latest = find_latest_hub_checkpoint(repo_id)
+    if latest is None:
+        raise FileNotFoundError(
+            f"No checkpoint found in '{repo_id}' under '{CHECKPOINTS_DIR}/'. "
+            "Was the run trained with --save_checkpoint_to_hub?"
+        )
+    snapshot_download(
+        repo_id=repo_id,
+        repo_type="model",
+        allow_patterns=f"{latest}/*",
+        local_dir=str(output_dir),
+    )
+    checkpoint_dir = output_dir / latest
+    update_last_checkpoint(checkpoint_dir)
+    return checkpoint_dir

src/lerobot/configs/__init__.py

@@ -22,7 +22,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
 """
 
 from .dataset import DatasetRecordConfig
-from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .recipe import MessageTurn, TrainingRecipe, load_recipe
 from .types import (
@@ -55,6 +55,7 @@ __all__ = [
     "DatasetRecordConfig",
     "DatasetConfig",
     "EvalConfig",
+    "JobConfig",
     "MessageTurn",
     "PeftConfig",
     "PreTrainedConfig",

src/lerobot/configs/default.py

@@ -145,3 +145,35 @@ class PeftConfig:
     # If None, the PEFT library defaults to alpha=8, which may dampen high-rank adapters.
     # Common values are r (alpha == rank) or 2*r.
     lora_alpha: int | None = None
+
+
+@dataclass
+class JobConfig:
+    # Where training runs. None (omitted) or "local" runs on this machine.
+    # Any other value is an HF Jobs flavor and submits the run to HF Jobs.
+    # List available flavors + pricing with `hf jobs hardware` command.
+    target: str | None = None
+    # Runtime image for the remote job (ignored for local runs).
+    image: str = "huggingface/lerobot-gpu:latest"
+    # Max wall-clock for the remote job as an HF Jobs duration string (e.g. "2h").
+    # Defaults to "2d": We pass an explicit, generous cap instead. Set a smaller
+    # value to fail fast, or a larger one for long runs.
+    timeout: str | None = "2d"
+    # Submit and exit instead of streaming the job logs in the foreground.
+    detach: bool = False
+    # Extra tags attached to the HF job and to any dataset this run pushes to the
+    # Hub. A "lerobot" tag is always added; e.g. --job.tags '["lelab"]' adds more.
+    tags: list[str] = field(default_factory=list)
+
+    # Two entry points to the same predicate: the staticmethod tests a raw target string
+    # straight from argv (before any JobConfig exists, to decide dispatch early), while the
+    # property is the ergonomic accessor for code that already holds a config instance.
+    @staticmethod
+    def is_remote_target(target: str | None) -> bool:
+        """True when `target` names an HF Jobs flavor rather than a local run."""
+        return target not in (None, "local")
+
+    @property
+    def is_remote(self) -> bool:
+        """True when training should run on HF Jobs rather than this machine."""
+        return self.is_remote_target(self.target)

src/lerobot/configs/train.py

@@ -26,11 +26,12 @@ from huggingface_hub.errors import HfHubHTTPError
 
 from lerobot import envs
 from lerobot.optim import LRSchedulerConfig, OptimizerConfig
-from lerobot.utils.hub import HubMixin
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR
+from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint
 from lerobot.utils.sample_weighting import SampleWeightingConfig
 
 from . import parser
-from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .rewards import RewardModelConfig
 
@@ -83,10 +84,11 @@ class TrainPipelineConfig(HubMixin):
     # with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
     output_dir: Path | None = None
     job_name: str | None = None
-    # Set `resume` to true to resume a previous run. In order for this to work, you will need to make sure
-    # `dir` is the directory of an existing run with at least one checkpoint in it.
-    # Note that when resuming a run, the default behavior is to use the configuration from the checkpoint,
-    # regardless of what's provided with the training command at the time of resumption.
+    # Set `resume` to true to resume a previous run. Pass `--config_path` pointing at either a local
+    # checkpoint's train_config.json or a Hub repo id holding `checkpoints/<step>/` subtrees (the
+    # latest checkpoint is downloaded and resumed from). Note that when resuming, the default behavior
+    # is to use the configuration from the checkpoint, regardless of what's provided with the training
+    # command at the time of resumption (CLI `--*` flags still override).
     resume: bool = False
     # `seed` is used for training (eg: model initialization, dataset shuffling)
     # AND for the evaluation environments.
@@ -118,6 +120,13 @@ class TrainPipelineConfig(HubMixin):
     wandb: WandBConfig = field(default_factory=WandBConfig)
     peft: PeftConfig | None = None
 
+    # Where to run training (local default, or an HF Jobs flavor). See JobConfig.
+    job: JobConfig = field(default_factory=JobConfig)
+    # Push each saved checkpoint to the Hub (policy.repo_id) as it is written, not
+    # just the final model (useful to monitor progress mid-run). Optional; the
+    # final model is pushed regardless. Works the same locally and remotely.
+    save_checkpoint_to_hub: bool = False
+
     # Sample weighting configuration (e.g., for RA-BC training)
     sample_weighting: SampleWeightingConfig | None = None
 
@@ -137,10 +146,17 @@ class TrainPipelineConfig(HubMixin):
             return self.reward_model  # type: ignore[return-value]
         return self.policy  # type: ignore[return-value]
 
-    def validate(self) -> None:
-        # HACK: We parse again the cli args here to get the pretrained paths if there was some.
-        policy_path = parser.get_path_arg("policy")
+    def _resolve_pretrained_from_cli(self) -> None:
+        """Resolve the pretrained source passed on the CLI into a loaded config.
+
+        The pretrained paths (`--policy.path`, `--reward_model.path`) and
+        `--config_path` are only recoverable by re-reading the CLI args: draccus
+        has already consumed them by the time `validate()` runs, so they are not
+        reflected on `self`. Exactly one source applies, in priority order:
+        reward-model path, policy path, then resume.
+        """
         reward_model_path = parser.get_path_arg("reward_model")
+        policy_path = parser.get_path_arg("policy")
 
         if reward_model_path:
             cli_overrides = parser.get_cli_overrides("reward_model")
@@ -149,31 +165,54 @@ class TrainPipelineConfig(HubMixin):
             )
             self.reward_model.pretrained_path = str(Path(reward_model_path))
         elif policy_path:
-            yaml_overrides = parser.get_yaml_overrides("policy")
-            cli_overrides = parser.get_cli_overrides("policy") or []
-            self.policy = PreTrainedConfig.from_pretrained(
-                policy_path, cli_overrides=yaml_overrides + cli_overrides
-            )
+            overrides = parser.get_yaml_overrides("policy") + (parser.get_cli_overrides("policy") or [])
+            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=overrides)
             self.policy.pretrained_path = Path(policy_path)
         elif self.resume:
-            config_path = parser.parse_arg("config_path")
-            if not config_path:
-                raise ValueError(
-                    f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
-                )
+            self._resolve_resume_checkpoint()
 
-            if not Path(config_path).resolve().exists():
-                raise NotADirectoryError(
-                    f"{config_path=} is expected to be a local path. "
-                    "Resuming from the hub is not supported for now."
-                )
+    def _resolve_resume_checkpoint(self) -> None:
+        """Point the trainable config at the checkpoint named by `--config_path`.
 
+        `config_path` is either a local path (to a checkpoint's train_config.json or its
+        pretrained_model/ dir) or a Hub repo id. For a Hub repo, the latest checkpoint is downloaded
+        into a fresh local run dir and resumed from there. The download is skipped when dispatching to
+        an HF Job (`job.is_remote`): the pod performs it when it runs the resume locally, and
+        `submit_to_hf` resolves the source repo for the remote command.
+        """
+        config_path = parser.parse_arg("config_path")
+        if not config_path:
+            raise ValueError(
+                f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
+            )
+
+        if Path(config_path).resolve().exists():
             policy_dir = Path(config_path).parent
-            if self.policy is not None:
-                self.policy.pretrained_path = policy_dir
-            if self.reward_model is not None:
-                self.reward_model.pretrained_path = str(policy_dir)
             self.checkpoint_path = policy_dir.parent
+        elif self.job.is_remote:
+            return
+        else:
+            from lerobot.common.train_utils import resolve_resume_checkpoint
+
+            # `self.output_dir` was loaded from the checkpoint's config and points at the original
+            # run's (now-absent) local dir. Resume into a fresh local dir instead, unless the user
+            # passed --output_dir explicitly.
+            cli_output_dir = parser.parse_arg("output_dir")
+            if cli_output_dir:
+                self.output_dir = Path(cli_output_dir)
+            else:
+                now = dt.datetime.now()
+                self.output_dir = Path("outputs/train") / f"{now:%Y-%m-%d}/{now:%H-%M-%S}_resume"
+            self.checkpoint_path = resolve_resume_checkpoint(config_path, self.output_dir)
+            policy_dir = self.checkpoint_path / PRETRAINED_MODEL_DIR
+
+        if self.policy is not None:
+            self.policy.pretrained_path = policy_dir
+        if self.reward_model is not None:
+            self.reward_model.pretrained_path = str(policy_dir)
+
+    def validate(self) -> None:
+        self._resolve_pretrained_from_cli()
 
         if self.policy is None and self.reward_model is None:
             raise ValueError(
@@ -216,9 +255,19 @@ class TrainPipelineConfig(HubMixin):
         if self.eval_steps > 0 and self.dataset.eval_split == 0.0:
             raise ValueError("eval_steps > 0 requires dataset.eval_split > 0.0 to hold out eval data.")
 
-        if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id:
+        # Remote runs auto-generate the repo_id in submit_to_hf (the policy may only be
+        # resolved here, from --policy.path), so don't demand it up front for them.
+        if (
+            hasattr(active_cfg, "push_to_hub")
+            and active_cfg.push_to_hub
+            and not active_cfg.repo_id
+            and not self.job.is_remote
+        ):
             raise ValueError("'repo_id' argument missing. Please specify it to push the model to the hub.")
 
+        if self.save_checkpoint_to_hub and not (self.policy is not None and self.policy.repo_id):
+            raise ValueError("save_checkpoint_to_hub requires --policy.repo_id.")
+
     @classmethod
     def __get_path_fields__(cls) -> list[str]:
         """Keys for draccus pretrained-path loading."""
@@ -255,22 +304,30 @@ class TrainPipelineConfig(HubMixin):
         elif Path(model_id).is_file():
             config_file = model_id
         else:
+            dl_kwargs = {
+                "repo_id": model_id,
+                "revision": revision,
+                "cache_dir": cache_dir,
+                "force_download": force_download,
+                "proxies": proxies,
+                "resume_download": resume_download,
+                "token": token,
+                "local_files_only": local_files_only,
+            }
             try:
-                config_file = hf_hub_download(
-                    repo_id=model_id,
-                    filename=TRAIN_CONFIG_NAME,
-                    revision=revision,
-                    cache_dir=cache_dir,
-                    force_download=force_download,
-                    proxies=proxies,
-                    resume_download=resume_download,
-                    token=token,
-                    local_files_only=local_files_only,
-                )
+                config_file = hf_hub_download(filename=TRAIN_CONFIG_NAME, **dl_kwargs)
             except HfHubHTTPError as e:
-                raise FileNotFoundError(
-                    f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
-                ) from e
+                # No root train_config.json: this is a repo of periodic checkpoints from an
+                # interrupted run. Fall back to the latest checkpoint's config so the run can be
+                # resumed straight from the repo with `--config_path=<repo>`.
+                latest = find_latest_hub_checkpoint(model_id, token=token, revision=revision)
+                if latest is None:
+                    raise FileNotFoundError(
+                        f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
+                    ) from e
+                config_file = hf_hub_download(
+                    filename=f"{latest}/{PRETRAINED_MODEL_DIR}/{TRAIN_CONFIG_NAME}", **dl_kwargs
+                )
 
         cli_args = kwargs.pop("cli_args", [])
         # Legacy RA-BC migration only applies to framework-saved checkpoints (always JSON).

src/lerobot/jobs/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from lerobot.utils.import_utils import require_package
+
+# LeRobotDataset (imported at module top in dataset.py) pulls in heavy dataset deps;
+# guard the optional dependency here so importing this package fails loudly if it's missing.
+require_package("datasets", extra="dataset")
+
+from .hf import submit_to_hf
+
+__all__ = ["submit_to_hf"]

src/lerobot/jobs/dataset.py

@@ -0,0 +1,53 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Make a training dataset reachable from an HF Job pod.
+
+The pod can't see the host's ~/.cache/huggingface/lerobot, so the dataset has to
+live on the Hub: the pod downloads it by repo_id at train time (the forwarded
+HF_TOKEN covers private datasets). A dataset already on the Hub is used as-is; a
+local-only dataset is pushed to a PRIVATE repo first (never public).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from lerobot.datasets import LeRobotDataset
+from lerobot.utils.constants import HF_LEROBOT_HOME
+
+if TYPE_CHECKING:
+    from huggingface_hub import HfApi
+
+
+def ensure_dataset_available(repo_id: str, *, api: HfApi, tags: list[str] | None = None) -> None:
+    """Ensure repo_id resolves on the Hub, pushing a local-only dataset privately first.
+
+    `tags` are attached to the dataset only when we push it (an already-on-Hub
+    dataset is left untouched). Raises RuntimeError if the dataset is neither on
+    the Hub nor in the local cache.
+    """
+    if api.repo_exists(repo_id, repo_type="dataset"):
+        return
+
+    local_present = (HF_LEROBOT_HOME / repo_id / "meta" / "info.json").is_file()
+    if not local_present:
+        raise RuntimeError(
+            f"Dataset '{repo_id}' is not in the local cache ({HF_LEROBOT_HOME}) and could not be "
+            f"reached on the Hub — it may not exist, or be private and inaccessible with your "
+            f"token. Record or download it first, or run `hf auth login`."
+        )
+
+    print(f"[dataset] '{repo_id}' is local-only; pushing to a PRIVATE Hub repo...")
+    LeRobotDataset(repo_id).push_to_hub(private=True, tags=tags)
+    print(f"[dataset] '{repo_id}' uploaded (private). The job will download it by repo_id.")

src/lerobot/jobs/hf.py

@@ -0,0 +1,425 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Run a lerobot training on HF Jobs (HuggingFace GPUs).
+
+Ported and simplified from lelab's runners/hf_cloud.py: no UI log queue, no
+registry — just submit and stream to stdout.
+"""
+
+from __future__ import annotations
+
+import copy
+import datetime as dt
+import json
+import netrc
+import os
+import re
+import signal
+import sys
+import tempfile
+import threading
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import httpx
+from huggingface_hub import (
+    HfApi,
+    create_repo,
+    fetch_job_logs,
+    get_token,
+    inspect_job,
+    run_job,
+    upload_file,
+)
+
+from lerobot.common.train_utils import push_checkpoint_to_hub
+from lerobot.configs import parser
+
+from .dataset import ensure_dataset_available
+
+if TYPE_CHECKING:
+    from lerobot.configs.train import TrainPipelineConfig
+
+_SLUG_RE = re.compile(r"[^a-zA-Z0-9._-]+")
+
+_TERMINAL_STAGES = {"COMPLETED", "CANCELED", "ERROR", "DELETED"}
+
+# huggingface_hub 1.x runs on httpx: transient HTTP/transport failures surface as
+# httpx.HTTPError and socket-level errors as OSError. Catching only these keeps real
+# bugs (TypeError, AttributeError, ...) from being silently retried or counted as
+# job failures.
+_TRANSIENT_NET_ERRORS = (OSError, httpx.HTTPError)
+
+# Always attached to remote jobs and pushed datasets so LeRobot-originated work
+# is identifiable on the Hub; callers (e.g. LeLab) add their own via --job.tags.
+LEROBOT_TAG = "lerobot"
+
+
+def resolve_job_tags(extra: list[str] | None) -> list[str]:
+    """Return the tag list for a run: the lerobot tag plus any extras, deduped, order-stable."""
+    tags = [LEROBOT_TAG, *(extra or [])]
+    seen: set[str] = set()
+    return [t for t in tags if not (t in seen or seen.add(t))]
+
+
+def resolve_wandb_api_key() -> str | None:
+    """Host's wandb key for forwarding to the job: $WANDB_API_KEY, else ~/.netrc."""
+    key = os.environ.get("WANDB_API_KEY")
+    if key:
+        return key
+    try:
+        rc = netrc.netrc()
+    except (FileNotFoundError, netrc.NetrcParseError, OSError):
+        return None
+    auth = rc.authenticators("api.wandb.ai")
+    if auth is None:
+        return None
+    _login, _account, password = auth
+    return password or None
+
+
+def build_repo_id(username: str, job_name: str, now: dt.datetime) -> str:
+    """Generate the model repo id for a remote run: <user>/<job_name>_<timestamp>."""
+    slug = _SLUG_RE.sub("-", job_name).strip("-") or "train"
+    stamp = now.strftime("%Y-%m-%d_%H-%M-%S")
+    return f"{username}/{slug}_{stamp}"
+
+
+def build_remote_config_file(cfg, repo_id: str, dest: Path, tags: list[str] | None = None) -> Path:
+    """Write a train_config.json for the pod, with remote overrides applied.
+
+    The pod runs `lerobot-train --config_path=<dest>` and downloads the dataset
+    by repo_id into its own cache. Client-only fields are stripped so the config
+    is accepted by the trainer image: `job` (pure client orchestration) is always
+    removed, and `save_checkpoint_to_hub` is removed unless explicitly enabled —
+    older lerobot images reject unknown keys, so the default keeps the config
+    compatible with the released `lerobot-gpu` image. `tags` are merged into
+    policy.tags so the trained model the pod pushes carries them too.
+    """
+    remote = copy.deepcopy(cfg)
+    remote.policy.push_to_hub = True
+    remote.policy.repo_id = repo_id
+    # Don't pin the client's resolved device (e.g. "mps"); let the pod auto-detect its GPU.
+    remote.policy.device = None
+    # Drop any host-local dataset root; the pod resolves the dataset by repo_id.
+    remote.dataset.root = None
+    if tags:
+        existing = list(remote.policy.tags or [])
+        remote.policy.tags = existing + [t for t in tags if t not in existing]
+
+    # Encode to the canonical, pod-parseable dict, then drop the keys the released
+    # trainer image doesn't know about.
+    data = remote.to_dict()
+    data.pop("job", None)
+    if not remote.save_checkpoint_to_hub:
+        data.pop("save_checkpoint_to_hub", None)
+
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    dest.write_text(json.dumps(data, indent=4))
+    return dest
+
+
+def _stage_config_on_hub(cfg, repo_id: str, token: str, tags: list[str] | None = None) -> str:
+    """Upload train_config.json to the model repo and return the repo_id for --config_path."""
+    create_repo(repo_id, repo_type="model", private=True, exist_ok=True, token=token)
+    with tempfile.TemporaryDirectory() as tmp:
+        config_path = build_remote_config_file(cfg, repo_id, Path(tmp) / "train_config.json", tags=tags)
+        upload_file(
+            path_or_fileobj=config_path,
+            path_in_repo="train_config.json",
+            repo_id=repo_id,
+            repo_type="model",
+            token=token,
+        )
+    return repo_id
+
+
+def _tail_logs(
+    job_id: str,
+    done: threading.Event,
+    success_marker: str | None = None,
+    success_event: threading.Event | None = None,
+) -> None:
+    """Stream job logs to stdout, reconnecting on dropped streams until done is set.
+
+    Each reconnect re-fetches the full buffered log, so we track how many lines
+    were already printed and skip them — otherwise a fast-failing job's traceback
+    gets reprinted on every reconnect.
+
+    When `success_marker` appears in a line, set `success_event` and `done` so the
+    caller can finish as soon as the trained model lands on the Hub, rather than
+    waiting out the platform's post-run finalization (which can add ~30s).
+    """
+    printed = 0
+    while not done.is_set():
+        try:
+            seen = 0
+            for line in fetch_job_logs(job_id=job_id, follow=True):
+                seen += 1
+                if seen <= printed:
+                    continue  # already shown on a previous connection
+                printed = seen
+                # fetch_job_logs yields SSE data without trailing newlines, so add one
+                # per entry — otherwise all log lines concatenate onto a single line.
+                print(line.rstrip("\n"), flush=True)
+                if success_marker and success_event is not None and success_marker in line:
+                    success_event.set()
+                    done.set()
+                    return
+                if done.is_set():
+                    return
+            # Stream closed cleanly. Wait a moment so the status poller can mark
+            # the job terminal before we reconnect (avoids re-tailing the buffer).
+            if done.wait(3):
+                return
+        except _TRANSIENT_NET_ERRORS:
+            if done.wait(2):
+                return
+
+
+def _poll_until_done(
+    job_id: str,
+    done: threading.Event,
+    poll_interval: float = 5.0,
+    status_holder: dict | None = None,
+    max_failures: int = 6,
+) -> str | None:
+    """Poll inspect_job until a terminal stage or until `done` is set.
+
+    Returns the terminal stage string, or None if `done` was set first (detach)
+    or after `max_failures` consecutive inspect_job errors. When a terminal stage
+    is reached and `status_holder` is given, records `status_holder["message"]`
+    (the platform's status message, e.g. "Job timeout").
+    """
+    failures = 0
+    while not done.is_set():
+        try:
+            info = inspect_job(job_id=job_id)
+            failures = 0
+            # `stage` is an enum in some huggingface_hub versions and a plain str in others.
+            stage = getattr(info.status.stage, "value", info.status.stage)
+            if stage in _TERMINAL_STAGES:
+                if status_holder is not None:
+                    status_holder["message"] = getattr(info.status, "message", None)
+                done.set()
+                return stage
+        except _TRANSIENT_NET_ERRORS:
+            failures += 1
+            if failures >= max_failures:
+                done.set()
+                return None
+        done.wait(poll_interval)
+    return None
+
+
+def _pod_forwarded_args(
+    argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
+) -> list[str]:
+    """User CLI overrides to replay on the pod, minus flags the submitter sets itself.
+
+    Handles both `--name=value` and `--name value` forms. Forwarding the user's overrides (e.g.
+    `--steps`, `--save_checkpoint_to_hub`) makes a remote resume behave like the same local command.
+    """
+    out: list[str] = []
+    skip_next = False
+    for i, tok in enumerate(argv):
+        if skip_next:
+            skip_next = False
+            continue
+        name = tok.split("=", 1)[0]
+        if name in drop_names or any(name.startswith(p) for p in drop_prefixes):
+            if "=" not in tok and i + 1 < len(argv) and not argv[i + 1].startswith("--"):
+                skip_next = True  # also drop the space-separated value
+            continue
+        out.append(tok)
+    return out
+
+
+def _build_resume_job(cfg: TrainPipelineConfig, username: str) -> tuple[str, list[str]]:
+    """Resolve the model repo and pod command to resume a run on a job.
+
+    A Hub `config_path` is resumed from directly: its checkpoint config already targets that repo,
+    so new checkpoints continue the lineage there. A local `config_path` has its checkpoint uploaded
+    to a new PRIVATE repo first, and the resumed run is forced to push back to it. The pod command
+    always carries `--job.target=local` so the checkpoint's saved `job.target` can't make the pod
+    re-dispatch itself.
+    """
+    config_path = parser.parse_arg("config_path")
+    forwarded = _pod_forwarded_args(
+        sys.argv[1:],
+        drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"),
+        drop_prefixes=("--job.",),
+    )
+
+    if Path(config_path).exists():
+        # Local checkpoint: stage it on the Hub so the pod can resume from it, and push back there.
+        # Resolve so a `last` symlink uploads under its real step name (digit), which the pod's
+        # latest-checkpoint lookup keys on.
+        checkpoint_dir = Path(cfg.checkpoint_path).resolve()
+        source_repo = build_repo_id(username, cfg.job_name or "train", dt.datetime.now(dt.UTC))
+        push_checkpoint_to_hub(checkpoint_dir, source_repo, private=True)
+        extra = [f"--policy.repo_id={source_repo}", "--policy.push_to_hub=true"]
+    else:
+        source_repo = config_path
+        extra = []
+
+    command = [
+        "lerobot-train",
+        *forwarded,
+        f"--config_path={source_repo}",
+        "--job.target=local",
+        *extra,
+    ]
+    return source_repo, command
+
+
+def submit_to_hf(cfg: TrainPipelineConfig) -> None:
+    """Submit a training job to HF Jobs infrastructure.
+
+    Validates cfg, resolves credentials, ensures the dataset is on the Hub, then either stages a
+    sanitized config (fresh run) or resumes from a checkpoint repo, submits the job, and tails logs
+    until completion or detaches immediately. Ctrl-C detaches without cancelling the remote job.
+    """
+    token = get_token()
+    if not token:
+        raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.")
+
+    api = HfApi(token=token)
+    user_info = api.whoami(token=token)
+    username = user_info["name"]
+
+    now = dt.datetime.now(dt.UTC)
+    fresh_repo_id: str | None = None
+    if not cfg.resume:
+        # Resolve the model repo and mark it for push BEFORE validate(): validate() requires repo_id
+        # to be set whenever push_to_hub is True. (A resume reuses the checkpoint's repo instead.)
+        if cfg.policy is not None:
+            base_name = cfg.job_name or cfg.policy.type
+            fresh_repo_id = cfg.policy.repo_id or build_repo_id(username, base_name, now)
+            cfg.policy.repo_id = fresh_repo_id
+            cfg.policy.push_to_hub = True
+        else:
+            # Path-based policy is resolved inside validate(); fall back to a generic slug.
+            fresh_repo_id = build_repo_id(username, cfg.job_name or "train", now)
+
+    cfg.validate()
+
+    if cfg.is_reward_model_training:
+        raise ValueError(
+            "Remote training via --job.target only supports policy training, not reward models. "
+            "Run reward-model training locally."
+        )
+
+    secrets: dict[str, str] = {"HF_TOKEN": token}
+    if cfg.wandb.enable:
+        wandb_key = resolve_wandb_api_key()
+        if wandb_key is None:
+            raise ValueError(
+                "wandb is enabled but no WANDB_API_KEY found. "
+                "Set it via `export WANDB_API_KEY=...` or add it to ~/.netrc."
+            )
+        secrets["WANDB_API_KEY"] = wandb_key
+
+    tags = resolve_job_tags(cfg.job.tags)
+    # The dataset must be reachable from the pod for both fresh and resumed runs; a local-only
+    # dataset is pushed PRIVATE here. Hoisted before the resume/fresh branch since it applies to both.
+    ensure_dataset_available(cfg.dataset.repo_id, api=api, tags=tags)
+
+    if cfg.resume:
+        repo_id, command = _build_resume_job(cfg, username)
+    else:
+        config_repo_id = _stage_config_on_hub(cfg, fresh_repo_id, token, tags=tags)
+        repo_id = fresh_repo_id
+        command = ["lerobot-train", f"--config_path={config_repo_id}"]
+
+    print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...")
+    job_info = run_job(
+        image=cfg.job.image,
+        command=command,
+        flavor=cfg.job.target,
+        secrets=secrets,
+        timeout=cfg.job.timeout,
+        # HF Jobs labels are key/value; expose each tag as a queryable label.
+        labels=dict.fromkeys(tags, "true"),
+    )
+    job_id = job_info.id
+    job_url = getattr(job_info, "url", None)
+    print(f"Job submitted: {job_id}")
+    if job_url:
+        print(f"  Job page:   {job_url}")
+    print(f"  Model repo: https://huggingface.co/{repo_id}")
+    print(f"  Monitor:    hf jobs logs {job_id}")
+    print(f"  Cancel:     hf jobs cancel {job_id}")
+
+    if cfg.job.detach:
+        return
+
+    done = threading.Event()
+    detached = threading.Event()
+    pushed_ok = threading.Event()
+    stage_holder: dict[str, str | None] = {}
+
+    def _poll() -> None:
+        stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
+
+    poll_thread = threading.Thread(target=_poll, daemon=True)
+    poll_thread.start()
+    # Finish as soon as the model is pushed, rather than waiting out the platform's
+    # post-run finalization before the job stage flips to COMPLETED. This matches the
+    # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay
+    # in sync. If it ever stops matching we just fall back to stage-based completion
+    # (~30s slower), so the contract is an optimization, not a correctness requirement.
+    success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
+    log_thread = threading.Thread(
+        target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True
+    )
+    log_thread.start()
+
+    def _detach(sig, frame):
+        detached.set()
+        done.set()
+        print("\nDetached. Job is still running.")
+        print(f"  Monitor: hf jobs logs {job_id}")
+        print(f"  Cancel:  hf jobs cancel {job_id}")
+
+    # signal.signal only works on the main thread; when called from a worker thread
+    # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
+    # handler rather than crashing with ValueError.
+    install_sigint = threading.current_thread() is threading.main_thread()
+    original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
+    if install_sigint:
+        signal.signal(signal.SIGINT, _detach)
+    try:
+        # Timeout-based join so SIGINT is delivered to the main thread promptly.
+        while poll_thread.is_alive():
+            poll_thread.join(timeout=0.5)
+        log_thread.join(timeout=5)
+    finally:
+        if install_sigint:
+            signal.signal(signal.SIGINT, original_sigint)
+
+    if detached.is_set():
+        return
+
+    if pushed_ok.is_set():
+        print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}")
+        return
+
+    stage = stage_holder.get("stage")
+    if stage != "COMPLETED":
+        message = stage_holder.get("message")
+        detail = f" ({message})" if message else ""
+        raise RuntimeError(
+            f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
+        )

src/lerobot/policies/molmoact2/configuration_molmoact2.py

@@ -79,6 +79,15 @@ class MolmoAct2Config(PreTrainedConfig):
     eval_seed: int | None = None
     rtc_config: RTCConfig | None = None
 
+    # Joint frame transform for cross-calibration compatibility.
+    # Some MolmoAct2 checkpoints were trained on data using a different joint
+    # convention than the current LeRobot calibration. Set both to apply a
+    # sign/offset correction at runtime (state before model, action after).
+    # See: https://huggingface.co/docs/lerobot/backwardcomp
+    # Default is None (no transform). Both must be set together.
+    joint_signs: list[float] | None = None
+    joint_offsets: list[float] | None = None
+
     # Default is full finetuning with gradients from the action expert flowing into the VLM.
     enable_lora_vlm: bool = False
     lora_rank: int = 64
@@ -123,6 +132,10 @@ class MolmoAct2Config(PreTrainedConfig):
 
     def __post_init__(self) -> None:
         super().__post_init__()
+        if (self.joint_signs is None) != (self.joint_offsets is None):
+            raise ValueError("joint_signs and joint_offsets must both be set or both be None.")
+        if self.joint_signs is not None and len(self.joint_signs) != len(self.joint_offsets):
+            raise ValueError("joint_signs and joint_offsets must have the same length.")
         if self.action_mode not in {"continuous", "discrete", "both"}:
             raise ValueError(
                 f"Unsupported action_mode={self.action_mode!r}. "

src/lerobot/policies/molmoact2/processor_molmoact2.py

@@ -1005,6 +1005,93 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep):
         return features
 
 
+@ProcessorStepRegistry.register(name="molmoact2_state_frame_transform")
+@dataclass
+class MolmoAct2StateFrameTransformStep(ProcessorStep):
+    """Convert robot state from arm frame to model frame before normalization.
+
+    Required for zero-shot deployment of MolmoAct2-SO100_101 on SO-100/101
+    arms calibrated with LeRobot >= 0.5.0 (v3.0 convention). The checkpoint
+    was trained on data using a different joint convention (sign flip on
+    shoulder_lift, 90 deg offset on shoulder_lift and elbow_flex).
+
+    No-op when joint_signs and joint_offsets are None (default), so this
+    step has no effect on fine-tuned models or other embodiments.
+
+    state_model = signs * arm_state + offsets
+
+    See: https://huggingface.co/docs/lerobot/backwardcomp
+    """
+
+    joint_signs: list[float] | None = None
+    joint_offsets: list[float] | None = None
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        if self.joint_signs is None or self.joint_offsets is None:
+            return transition
+        observation = transition.get(TransitionKey.OBSERVATION)
+        if not isinstance(observation, dict) or OBS_STATE not in observation:
+            return transition
+        transition = transition.copy()
+        observation = observation.copy()
+        state = torch.as_tensor(observation[OBS_STATE], dtype=torch.float32).clone()
+        n = len(self.joint_signs)
+        signs = torch.tensor(self.joint_signs, dtype=torch.float32, device=state.device)
+        offsets = torch.tensor(self.joint_offsets, dtype=torch.float32, device=state.device)
+        state[..., :n] = signs * state[..., :n] + offsets
+        observation[OBS_STATE] = state
+        transition[TransitionKey.OBSERVATION] = observation
+        return transition
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+    def get_config(self) -> dict[str, Any]:
+        return {"joint_signs": self.joint_signs, "joint_offsets": self.joint_offsets}
+
+
+@ProcessorStepRegistry.register(name="molmoact2_action_frame_transform")
+@dataclass
+class MolmoAct2ActionFrameTransformStep(ProcessorStep):
+    """Convert model action from model frame back to arm frame after unnormalization.
+
+    Inverse of MolmoAct2StateFrameTransformStep. Required for zero-shot
+    MolmoAct2-SO100_101 on SO-100/101 arms. No-op when both fields are None.
+
+    action_arm = signs * (model_action - offsets)
+
+    See: https://huggingface.co/docs/lerobot/backwardcomp
+    """
+
+    joint_signs: list[float] | None = None
+    joint_offsets: list[float] | None = None
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        if self.joint_signs is None or self.joint_offsets is None:
+            return transition
+        action = transition.get(TransitionKey.ACTION)
+        if action is None:
+            return transition
+        transition = transition.copy()
+        action = torch.as_tensor(action, dtype=torch.float32).clone()
+        n = len(self.joint_signs)
+        signs = torch.tensor(self.joint_signs, dtype=torch.float32, device=action.device)
+        offsets = torch.tensor(self.joint_offsets, dtype=torch.float32, device=action.device)
+        action[..., :n] = signs * (action[..., :n] - offsets)
+        transition[TransitionKey.ACTION] = action
+        return transition
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+    def get_config(self) -> dict[str, Any]:
+        return {"joint_signs": self.joint_signs, "joint_offsets": self.joint_offsets}
+
+
 @ProcessorStepRegistry.register(name="molmoact2_clamp_action")
 @dataclass
 class MolmoAct2ClampActionProcessorStep(ProcessorStep):
@@ -1067,6 +1154,10 @@ def make_molmoact2_pre_post_processors(
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),
         AddBatchDimensionProcessorStep(),
+        MolmoAct2StateFrameTransformStep(
+            joint_signs=config.joint_signs,
+            joint_offsets=config.joint_offsets,
+        ),
         MolmoAct2MaskedNormalizerProcessorStep(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
@@ -1102,6 +1193,10 @@ def make_molmoact2_pre_post_processors(
             norm_map=config.normalization_mapping,
             stats=masked_dataset_stats,
         ),
+        MolmoAct2ActionFrameTransformStep(
+            joint_signs=config.joint_signs,
+            joint_offsets=config.joint_offsets,
+        ),
         DeviceProcessorStep(device="cpu"),
     ]
 

src/lerobot/policies/pretrained.py

@@ -340,6 +340,9 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
                 ignore_patterns=["*.tmp", "*.log"],
             )
 
+            # Contract: lerobot.jobs.hf.submit_to_hf watches for this exact
+            # "Model pushed to <url>" line to end a remote run early. Keep the wording
+            # and URL format in sync (it falls back to status polling if they drift).
             logging.info(f"Model pushed to {commit_info.repo_url.url}")
 
     def generate_model_card(

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -17,12 +17,10 @@ from __future__ import annotations
 import logging
 from collections import deque
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-import numpy as np
 import torch
 import torch.nn.functional as F  # noqa: N812
-from PIL import Image
 from torch import Tensor, nn
 
 from lerobot.policies.pretrained import PreTrainedPolicy, T
@@ -55,12 +53,13 @@ class VLAJEPAModel(nn.Module):
       - DiT-B: flow-matching action head for future action prediction
       - V-JEPA: world model for video frame prediction
 
-    Input: List[dict] native format (same as original starVLA)
-      - "image": List[PIL.Image] (multi-view images)
-      - "video": np.ndarray [V, T, H, W, 3]
-      - "lang": str (task instruction)
-      - "action": np.ndarray [T, action_dim] (optional, training only)
-      - "state": np.ndarray [1, state_dim] (optional)
+    Inputs are batched tensors kept on the model device
+      - images: List[List[Tensor [C, H, W]]] (float [0,1]) — per sample, per view (Qwen messages)
+      - instructions: List[str]
+      - videos: Tensor [B, V, T, C, H, W] (float [0,1], world model only)
+      - actions: Tensor [B, T, action_dim] (optional, training only)
+      - state: Tensor [B, 1, state_dim] (optional)
+      - action_is_pad: Tensor [B, T] (optional)
     """
 
     def __init__(self, config: VLAJEPAConfig) -> None:
@@ -75,6 +74,11 @@ class VLAJEPAModel(nn.Module):
         self.action_tokens, self.action_token_ids, self.embodied_action_token_id = (
             self.qwen.expand_tokenizer()
         )
+        self.register_buffer(
+            "_action_token_ids_t",
+            torch.tensor(self.action_token_ids, dtype=torch.long),
+            persistent=False,
+        )
 
         # Action head (flow-matching DiT)
         self.action_model = VLAJEPAActionHead(config, cross_attention_dim=self.qwen.model.config.hidden_size)
@@ -161,166 +165,123 @@ class VLAJEPAModel(nn.Module):
 
     # ---- Native VLA-JEPA forward (follows original VLA_JEPA.py) ----
 
-    def forward(self, examples: list[dict]) -> dict[str, Tensor]:
-        """
-        Native forward pass following original starVLA VLA_JEPA.forward.
-
-        Args:
-            examples: List of per-sample dicts with keys:
-                "image"  : List[PIL.Image]  — multi-view images
-                "video"  : np.ndarray [V, T, H, W, 3]
-                "lang"   : str — task instruction
-                "action" : np.ndarray [T, action_dim] (optional)
-                "state"  : np.ndarray [1, state_dim] (optional)
-
-        Returns:
-            dict with "action_loss" and "wm_loss" keys (scalar Tensors).
-        """
-        # Unpack native format (same pattern as original VLA_JEPA.py)
-        batch_images = [ex["image"] for ex in examples]  # List[List[PIL.Image]]
-        batch_videos = [ex["video"] for ex in examples]  # List[np.ndarray]
-        instructions = [ex["lang"] for ex in examples]  # List[str]
-        has_action = "action" in examples[0] and examples[0]["action"] is not None
-        actions = [ex["action"] for ex in examples] if has_action else None
-        has_state = "state" in examples[0] and examples[0]["state"] is not None
-        state = [ex["state"] for ex in examples] if has_state else None
-        action_is_pad = (
-            [ex["action_is_pad"] for ex in examples]
-            if has_action and "action_is_pad" in examples[0] and examples[0]["action_is_pad"] is not None
-            else None
-        )
-
-        # Stack videos: [B, V, T, H, W, 3] -> [B, V, T, 3, H, W]
-        batch_videos = np.stack(batch_videos)
-        batch_videos = batch_videos.transpose(0, 1, 2, 5, 3, 4)  # [B, V, T, 3, H, W]
-
-        # Adjust number of views for the world model:
-        # - fewer views than expected: duplicate the first view to fill up
-        # - more views than expected: keep only the first num_views_world_model views
-        num_views_world_model = self.config.jepa_tubelet_size
-        if batch_videos.shape[1] < num_views_world_model:
-            num_missing_views = num_views_world_model - batch_videos.shape[1]
-            first_view = np.repeat(batch_videos[:, :1], num_missing_views, axis=1)
-            batch_videos = np.concatenate([batch_videos, first_view], axis=1)
-        elif batch_videos.shape[1] > num_views_world_model:
-            batch_videos = batch_videos[:, :num_views_world_model]
-
-        # ---- Step 1: QwenVL encode (same as original) ----
+    def _encode_qwen(
+        self, images: list[list[Tensor]], instructions: list[str], *, need_action_tokens: bool
+    ) -> tuple[Tensor, Tensor, Tensor | None]:
+        """Run Qwen and gather the embodied-action (and optionally action) token hidden states."""
         qwen_inputs = self.qwen.build_inputs(
-            images=batch_images,
+            images=images,
             instructions=instructions,
             action_prompt=self.replace_prompt,
             embodied_prompt=self.embodied_replace_prompt,
         )
-
-        # Locate embodied-action tokens (always needed for action head)
-        embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id
-        embodied_indices = embodied_mask.nonzero(as_tuple=True)
-
-        # Locate action tokens (only needed for world model predictor)
-        if self.config.enable_world_model:
-            action_mask = torch.isin(
-                qwen_inputs["input_ids"],
-                torch.tensor(self.action_token_ids, device=qwen_inputs["input_ids"].device),
-            )
-            action_indices = action_mask.nonzero(as_tuple=True)
+        input_ids = qwen_inputs["input_ids"]
+        embodied_idx = (input_ids == self.embodied_action_token_id).nonzero(as_tuple=True)
+        action_idx = None
+        if need_action_tokens:
+            action_mask = torch.isin(input_ids, self._action_token_ids_t)
+            action_idx = action_mask.nonzero(as_tuple=True)
 
         device_type = next(self.parameters()).device.type
-
         with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
             last_hidden = self._qwen_last_decoder_hidden(qwen_inputs)  # [B, seq_len, H]
             b, _, h = last_hidden.shape
+            embodied_action_tokens = last_hidden[embodied_idx[0], embodied_idx[1], :].view(b, -1, h)
+            action_tokens = (
+                last_hidden[action_idx[0], action_idx[1], :].view(b, -1, h)
+                if action_idx is not None
+                else None
+            )
+        return embodied_action_tokens, action_tokens
 
-            if self.config.enable_world_model:
-                action_tokens = last_hidden[action_indices[0], action_indices[1], :].view(b, -1, h)
+    def _world_model_loss(self, videos: Tensor, action_tokens: Tensor) -> Tensor:
+        """JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1]."""
+        # Match the world model's expected view count: pad with the first view, or trim extras.
+        num_views = self.config.jepa_tubelet_size
+        if videos.shape[1] < num_views:
+            missing = num_views - videos.shape[1]
+            videos = torch.cat([videos, videos[:, :1].repeat(1, missing, 1, 1, 1, 1)], dim=1)
+        elif videos.shape[1] > num_views:
+            videos = videos[:, :num_views]
 
-            embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h)
+        b, v, t_frames, c, h_img, w_img = videos.shape
+        flat = videos.reshape(b * v, t_frames, c, h_img, w_img)
+        # Fast (torchvision) video processor on-device, do_rescale=False (frames already in [0, 1]).
+        video_pixels = self.video_processor(
+            videos=list(flat),
+            return_tensors="pt",
+            device=self.video_encoder.device,
+            do_rescale=False,
+        )["pixel_values_videos"]  # [B*V, T, C, H, W]
 
-        # ---- Step 2+3: JEPA Encoder + Predictor ----
-        device_wm = last_hidden.device
-        if not self.config.enable_world_model:
-            wm_loss = torch.tensor(0.0, device=device_wm)
+        with torch.no_grad():
+            video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels)
+            # Merge views: [B*V, ...] -> [B, ..., V*embed_dim]
+            video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2)
+
+        tubelet_size = self.video_encoder.config.tubelet_size
+        # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression
+        t_enc_total = self.config.num_video_frames // tubelet_size
+        if t_enc_total < 2:
+            return torch.zeros((), device=video_embeddings.device)
+
+        # Shift-by-one JEPA split: input_states = positions 0..T-2, gt_states = positions 1..T-1
+        t_enc_ctx = t_enc_total - 1
+        tokens_per_frame = video_embeddings.shape[1] // t_enc_total
+        input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :]
+        gt_states = video_embeddings[:, tokens_per_frame:, :]
+
+        expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep
+        if action_tokens.shape[1] < expected_actions:
+            pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1)
+            action_tokens = torch.cat([action_tokens, pad], dim=1)
+
+        predicted_states = self.video_predictor(
+            input_states.float(), action_tokens[:, :expected_actions].float()
+        )
+        return F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
+
+    def _action_loss(
+        self,
+        embodied_action_tokens: Tensor,
+        actions: Tensor,
+        state: Tensor | None,
+        action_is_pad: Tensor | None,
+    ) -> Tensor:
+        """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`."""
+        device_type = next(self.parameters()).device.type
+        with torch.autocast(device_type=device_type, dtype=torch.float32):
+            r = self.config.repeated_diffusion_steps
+            horizon = self.config.chunk_size
+            actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1)
+            embodied = embodied_action_tokens.repeat(r, 1, 1)
+            state_rep = state.to(embodied_action_tokens.dtype).repeat(r, 1, 1) if state is not None else None
+            pad_rep = action_is_pad[:, -horizon:].repeat(r, 1) if action_is_pad is not None else None
+            return self.action_model(embodied, actions_target, state_rep, pad_rep)
+
+    def forward(
+        self,
+        images: list[list[Tensor]],
+        instructions: list[str],
+        videos: Tensor | None = None,
+        actions: Tensor | None = None,
+        state: Tensor | None = None,
+        action_is_pad: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        """Native forward: Qwen encode → optional world-model loss → optional action-head loss."""
+        embodied_action_tokens, action_tokens = self._encode_qwen(
+            images, instructions, need_action_tokens=self.config.enable_world_model
+        )
+
+        if self.config.enable_world_model and videos is not None:
+            wm_loss = self._world_model_loss(videos, action_tokens)
         else:
-            b, v, t_frames, c, h_img, w_img = batch_videos.shape
-            batch_videos_flat = batch_videos.reshape(b * v, t_frames, c, h_img, w_img)
+            wm_loss = torch.zeros((), device=embodied_action_tokens.device)
 
-            video_pixels = self.video_processor(videos=list(batch_videos_flat), return_tensors="pt")[
-                "pixel_values_videos"
-            ].to(self.video_encoder.device)  # [B*V, T, C, H, W]
-
-            with torch.no_grad():
-                video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels)
-                # Merge views: [B*V, ...] -> [B, ..., V*embed_dim]
-                video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2)
-
-            tubelet_size = self.video_encoder.config.tubelet_size
-            device_wm = video_embeddings.device
-            # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression
-            t_enc_total = self.config.num_video_frames // tubelet_size
-
-            if t_enc_total < 2:
-                wm_loss = torch.tensor(0.0, device=device_wm)
-            else:
-                # Shift-by-one JEPA split (matches original VLA_JEPA.py lines 231-232):
-                # input_states: positions 0..T-2, gt_states: positions 1..T-1
-                t_enc_ctx = t_enc_total - 1
-                tokens_per_frame = video_embeddings.shape[1] // t_enc_total
-
-                input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :]
-                gt_states = video_embeddings[:, tokens_per_frame:, :]
-
-                expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep
-                if action_tokens.shape[1] < expected_actions:
-                    pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1)
-                    action_tokens = torch.cat([action_tokens, pad], dim=1)
-
-                predicted_states = self.video_predictor(
-                    input_states.float(),
-                    action_tokens[:, :expected_actions].float(),
-                )
-
-                wm_loss = F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
-
-        if not has_action:
+        if actions is None:
             return {"wm_loss": wm_loss}
 
-        # ---- Step 4: Action Head ----
-        with torch.autocast(device_type=device_type, dtype=torch.float32):
-            actions_tensor = torch.tensor(
-                np.array(actions), device=last_hidden.device, dtype=torch.float32
-            )  # [B, T_full, action_dim]
-            action_horizon = self.config.chunk_size
-            actions_target = actions_tensor[:, -action_horizon:, :]
-
-            state_tensor = None
-            if state is not None:
-                state_tensor = torch.tensor(
-                    np.array(state), device=last_hidden.device, dtype=last_hidden.dtype
-                )  # [B, 1, state_dim]
-
-            repeated_diffusion_steps = self.config.repeated_diffusion_steps
-            actions_target = actions_target.repeat(repeated_diffusion_steps, 1, 1)
-            embodied_action_tokens = embodied_action_tokens.repeat(repeated_diffusion_steps, 1, 1)
-            if state_tensor is not None:
-                state_tensor = state_tensor.repeat(repeated_diffusion_steps, 1, 1)
-
-            action_is_pad_rep = None
-            if action_is_pad is not None:
-                pad_tensor = torch.stack(
-                    [
-                        p.to(actions_target.device)
-                        if isinstance(p, Tensor)
-                        else torch.tensor(p, device=actions_target.device)
-                        for p in action_is_pad
-                    ]
-                )  # [B, T_full]
-                pad_tensor = pad_tensor[:, -action_horizon:]  # [B, action_horizon]
-                action_is_pad_rep = pad_tensor.repeat(repeated_diffusion_steps, 1)  # [B*R, action_horizon]
-
-            action_loss = self.action_model(
-                embodied_action_tokens, actions_target, state_tensor, action_is_pad_rep
-            )
-
+        action_loss = self._action_loss(embodied_action_tokens, actions, state, action_is_pad)
         return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight}
 
     # ---- Native predict_action (follows original VLA_JEPA.predict_action) ----
@@ -328,58 +289,23 @@ class VLAJEPAModel(nn.Module):
     @torch.no_grad()
     def predict_action(
         self,
-        batch_images: list[list[Image.Image]],
+        images: list[list[Tensor]],
         instructions: list[str],
-        state: np.ndarray | None = None,
-    ) -> np.ndarray:
-        """
-        Native action prediction following original VLA_JEPA.predict_action.
-
-        Args:
-            batch_images: List of samples; each is List[PIL.Image] (multi-view).
-            instructions: Task instructions, one per sample.
-            state: Optional [B, state_dim] numpy array.
-
-        Returns:
-            np.ndarray [B, action_horizon, action_dim] — predicted actions.
-        """
+        state: Tensor | None = None,
+    ) -> Tensor:
+        """Predict an action chunk. `images` is per-sample, per-view float [0,1] [C, H, W] tensors."""
         if self.config.resize_images_to is not None:
             height, width = self.config.resize_images_to
-            resampling = getattr(Image, "Resampling", Image).BOX
-            batch_images = [
-                [image.resize((width, height), resample=resampling) for image in sample_images]
-                for sample_images in batch_images
+            images = [
+                [F.interpolate(img[None], size=(height, width), mode="area")[0] for img in views]
+                for views in images
             ]
 
-        qwen_inputs = self.qwen.build_inputs(
-            images=batch_images,
-            instructions=instructions,
-            action_prompt=self.replace_prompt,
-            embodied_prompt=self.embodied_replace_prompt,
+        embodied_action_tokens, _ = self._encode_qwen(images, instructions, need_action_tokens=False)
+        return self.action_model.predict_action(
+            embodied_action_tokens.float(), state.float() if state is not None else None
         )
 
-        embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id
-        embodied_indices = embodied_mask.nonzero(as_tuple=True)
-
-        device_type = next(self.parameters()).device.type
-
-        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
-            last_hidden = self._qwen_last_decoder_hidden(qwen_inputs)  # [B, seq_len, H]
-            b, _, h = last_hidden.shape
-            embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h)
-
-        state_tensor = None
-        if state is not None:
-            state_tensor = torch.from_numpy(np.array(state)).to(
-                device=last_hidden.device, dtype=last_hidden.dtype
-            )
-
-        pred_actions = self.action_model.predict_action(
-            embodied_action_tokens.float(), state_tensor.float() if state_tensor is not None else None
-        )  # [B, action_horizon, action_dim]
-
-        return pred_actions.detach().cpu().numpy()
-
 
 # ============================================================================
 # LeRobot Adapter Layer - converts between LeRobot batch format and native VLA-JEPA format
@@ -390,9 +316,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
     """
     LeRobot adapter for VLA-JEPA.
 
-    Converts LeRobot's standard batch format (dict[str, Tensor]) to the native
-    VLA-JEPA format (List[dict]), calls the native model, and converts outputs
-    back to LeRobot format.
+    Converts LeRobot's standard batch format (dict[str, Tensor]) to the batched tensors
+    the native model expects (keeping everything on-device), calls the native model, and
+    converts outputs back to LeRobot format.
     """
 
     config_class = VLAJEPAConfig
@@ -419,9 +345,8 @@ class VLAJEPAPolicy(PreTrainedPolicy):
 
     # ---- Format Conversion: LeRobot → Native ----
 
-    def _prepare_model_inputs(self, batch: dict[str, Tensor]) -> list[dict]:
-        """
-        Convert LeRobot batch format to native VLA-JEPA examples format.
+    def _prepare_model_inputs(self, batch: dict[str, Tensor], training=True) -> dict[str, Any]:
+        """Convert a LeRobot batch to the model's batched, on-device inputs.
 
         LeRobot format:
             batch = {
@@ -431,65 +356,25 @@ class VLAJEPAPolicy(PreTrainedPolicy):
                 "task": str | List[str],  (optional instruction)
             }
 
-        Native format (List[dict]):
-            {
-                "image": List[PIL.Image],       # multi-view images per sample
-                "video": np.ndarray [V, T, H, W, 3],
-                "lang": str,                     # task instruction
-                "action": np.ndarray [T, action_dim],  # optional
-                "state": np.ndarray [1, state_dim],    # optional
-            }
+        Returns the kwargs for `VLAJEPAModel.forward` / `.predict_action` (everything stays
+        on the batch device; no per-sample shredding): `images` (per-sample, per-view list for
+        Qwen messages), `instructions`, and the batched `videos` / `actions` / `state` /
+        `action_is_pad` when present.
         """
-        # Determine batch size from the first image feature
         image_keys = list(self.config.image_features.keys())
         if not image_keys:
             raise ValueError("VLAJEPA requires at least one image feature.")
-        first_key = image_keys[0]
-        first_tensor = batch[first_key]
-        batch_size = first_tensor.shape[0]
+        batch_size = batch[image_keys[0]].shape[0]
 
-        # ---- Collect images per sample ----
-        # images_per_sample[b][v] = PIL.Image for view v
-        images_per_sample: list[list[Image.Image]] = [[] for _ in range(batch_size)]
+        # Current-frame image per view ([B, C, H, W]); regroup per sample for Qwen messages.
+        frames = []
         for key in image_keys:
-            tensor = batch[key]  # [B, C, H, W] or [B, T, C, H, W]
-            if tensor.ndim == 5:
-                # observation_delta_indices = [0, 1, ..., num_video_frames-1]
-                # index 0 is the current observation (delta=0)
-                tensor = tensor[:, 0]
-            for b in range(batch_size):
-                images_per_sample[b].append(self.model.qwen.tensor_to_pil(tensor[b]))
+            t = batch[key]
+            if t.ndim == 5:  # [B, T, C, H, W] -> current observation (delta=0)
+                t = t[:, 0]
+            frames.append(self.model.qwen.to_pixel_values(t))
+        images = [[frame[b] for frame in frames] for b in range(batch_size)]
 
-        # ---- Collect videos per sample ----
-        # Build video arrays: for each sample, stack views as [V, T, H, W, 3]
-        # Check whether any image feature has a time dimension
-        video_source = None
-        for k in image_keys:
-            if k in batch:
-                video_source = batch[k]  # Use first available for shape inspection
-                break
-
-        if video_source is None:
-            raise ValueError("No image data found in batch for video construction.")
-
-        videos_per_sample = []
-        for b in range(batch_size):
-            sample_views = []
-            for k in image_keys:
-                t = batch[k][b]  # [C, H, W] or [T, C, H, W]
-                if t.ndim == 3:
-                    t = t.unsqueeze(0)  # [1, C, H, W]
-                # Convert to [T, H, W, 3] numpy
-                t_np = t.permute(0, 2, 3, 1).detach().cpu().float().numpy()
-                # Clamp to [0, 255]
-                if t_np.max() <= 1.0:
-                    t_np = t_np * 255.0
-                t_np = np.rint(t_np.clip(0, 255)).astype(np.uint8)
-                sample_views.append(t_np)
-            # Stack views: [V, T, H, W, 3]
-            videos_per_sample.append(np.stack(sample_views, axis=0))
-
-        # ---- Collect instructions ----
         tasks = batch.get("task")
         if tasks is None:
             instructions = ["Execute the robot action."] * batch_size
@@ -498,52 +383,32 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         else:
             instructions = list(tasks)
 
-        # ---- Collect actions (training only) ----
-        actions_list = None
-        action_is_pad_list = None
-        actions_tensor = batch.get(ACTION)
-        if actions_tensor is not None:
-            if actions_tensor.ndim == 2:
-                actions_tensor = actions_tensor.unsqueeze(1)
-            actions_list = [actions_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)]
-            action_is_pad_tensor = batch.get("action_is_pad")
-            if action_is_pad_tensor is not None:
-                action_is_pad_list = [action_is_pad_tensor[b].detach().cpu() for b in range(batch_size)]
+        inputs: dict[str, Any] = {"images": images, "instructions": instructions}
 
-        # ---- Collect state ----
-        state_list = None
-        state_tensor = batch.get(OBS_STATE)
-        if state_tensor is not None:
-            if state_tensor.ndim > 2:
-                state_tensor = state_tensor[:, -1, :]
-            if state_tensor.ndim == 2:
-                state_tensor = state_tensor.unsqueeze(1)  # [B, 1, state_dim]
-            state_list = [state_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)]
+        # Videos [B, V, T, C, H, W] - only assembled during training when the world model consumes them.
+        if self.model.config.enable_world_model and training:
+            views = [batch[k].unsqueeze(1) if batch[k].ndim == 4 else batch[k] for k in image_keys]
+            inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(views, dim=1))
 
-        # ---- Assemble native examples ----
-        examples = []
-        for b in range(batch_size):
-            example = {
-                "image": images_per_sample[b],
-                "video": videos_per_sample[b],
-                "lang": instructions[b],
-            }
-            if actions_list is not None:
-                example["action"] = actions_list[b]
-            if action_is_pad_list is not None:
-                example["action_is_pad"] = action_is_pad_list[b]
-            if state_list is not None:
-                example["state"] = state_list[b]
-            examples.append(example)
+        actions = batch.get(ACTION)
+        if actions is not None:
+            inputs["actions"] = (actions.unsqueeze(1) if actions.ndim == 2 else actions).float()
+            if (pad := batch.get("action_is_pad")) is not None:
+                inputs["action_is_pad"] = pad
 
-        return examples
+        state = batch.get(OBS_STATE)
+        if state is not None:
+            if state.ndim > 2:
+                state = state[:, -1, :]
+            inputs["state"] = (state.unsqueeze(1) if state.ndim == 2 else state).float()  # [B, 1, dim]
+
+        return inputs
 
     # ---- LeRobot Policy Interface ----
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """LeRobot train forward: convert → native forward → aggregate losses."""
-        examples = self._prepare_model_inputs(batch)
-        native_output = self.model.forward(examples)
+        native_output = self.model.forward(**self._prepare_model_inputs(batch, training=True))
 
         ref = next(iter(native_output.values()))
         zero = torch.zeros((), device=ref.device, dtype=ref.dtype)
@@ -561,16 +426,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         self.eval()
         self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
-        examples = self._prepare_model_inputs(batch)
-        batch_images = [ex["image"] for ex in examples]
-        instructions = [ex["lang"] for ex in examples]
-
-        state_np = None
-        if "state" in examples[0] and examples[0]["state"] is not None:
-            state_np = np.stack([ex["state"] for ex in examples])
-
-        actions_np = self.model.predict_action(batch_images, instructions, state_np)
-        return torch.from_numpy(actions_np).to(device=self.config.device, dtype=torch.float32)
+        inputs = self._prepare_model_inputs(batch, training=False)
+        actions = self.model.predict_action(inputs["images"], inputs["instructions"], inputs.get("state"))
+        return actions.to(device=self.config.device, dtype=torch.float32)
 
     @torch.no_grad()
     def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:

src/lerobot/policies/vla_jepa/qwen_interface.py

@@ -17,9 +17,7 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-import numpy as np
 import torch
-from PIL import Image
 
 from lerobot.utils.import_utils import _transformers_available
 
@@ -78,7 +76,7 @@ class Qwen3VLInterface(torch.nn.Module):
 
     def build_inputs(
         self,
-        images: Sequence[Sequence[Image.Image]],
+        images: Sequence[Sequence[torch.Tensor]],
         instructions: Sequence[str],
         action_prompt: str,
         embodied_prompt: str,
@@ -94,24 +92,42 @@ class Qwen3VLInterface(torch.nn.Module):
             content.append({"type": "text", "text": prompt})
             messages.append([{"role": "user", "content": content}])
 
+        # The Qwen image processor is a torchvision-backed fast processor: passing the
+        # images as GPU tensors (with `device`) keeps the whole vision pipeline on-device
+        # and avoids a GPU->CPU->GPU roundtrip. The image tensors are forwarded through
+        # apply_chat_template untouched into Qwen3VLProcessor.__call__.
+        # do_rescale=False: images already arrive as float in [0, 1] (the dataset decoder
+        # yields float32/255 and VISUAL normalization is IDENTITY), so we skip the
+        # processor's /255 rescale instead of round-tripping through uint8.
         batch_inputs = self.processor.apply_chat_template(
             messages,
             tokenize=True,
             add_generation_prompt=True,
             return_dict=True,
-            processor_kwargs={"padding": True, "return_tensors": "pt"},
+            processor_kwargs={
+                "padding": True,
+                "return_tensors": "pt",
+                "device": self.model.device,
+                "do_rescale": False,
+            },
         )
         return batch_inputs.to(self.model.device)
 
     @staticmethod
-    def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image:
-        image = image_tensor.detach().cpu()
-        if image.ndim == 3 and image.shape[0] in (1, 3):
-            image = image.permute(1, 2, 0)
-        image = image.float()
-        if image.max() <= 1.0:
-            image = image * 255.0
-        image = image.clamp(0, 255).round().to(torch.uint8).numpy()
-        if image.shape[-1] == 1:
-            image = np.repeat(image, 3, axis=-1)
-        return Image.fromarray(image)
+    def to_pixel_values(image_tensor: torch.Tensor) -> torch.Tensor:
+        """Prepare an image/video tensor for the fast processors (used with do_rescale=False).
+
+        The dataset decoder yields float32 in [0, 1] (channels-first) and VISUAL
+        normalization is IDENTITY, so the tensor already arrives in [0, 1]; we pass it
+        through as float and let the processors normalize (no rescale, no uint8
+        quantization). A single channel is expanded to 3 to match the RGB processors.
+
+        Works for any channels-first layout (channel dim is -3): [C, H, W], [B, C, H, W],
+        [T, C, H, W], [B, V, T, C, H, W], ...
+        """
+        image = image_tensor.detach().float()
+        if image.shape[-3] == 1:
+            repeats = [1] * image.ndim
+            repeats[-3] = 3
+            image = image.repeat(*repeats)
+        return image

src/lerobot/rollout/context.py

@@ -320,7 +320,9 @@ def build_rollout_context(
             raise ValueError(
                 f"Visual feature mismatch between policy and robot hardware.\n"
                 f"Policy expects: {expected_visuals}\n"
-                f"Robot provides: {provided_visuals}"
+                f"Robot provides: {provided_visuals}\n"
+                f"Use --rename_map to map camera names, e.g. "
+                f"""--rename_map='{{"observation.images.top": "observation.images.cam0"}}'"""
             )
 
     # --- 5. Dataset -------------

src/lerobot/scripts/lerobot_dataset_viz.py

@@ -77,6 +77,21 @@ from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD
 from lerobot.utils.utils import init_logging
 
 
+def get_feature_names(dataset: LeRobotDataset, key: str) -> list[str]:
+    """Return per-dimension names for a feature from the dataset metadata.
+
+    Only flat-list ``names`` metadata is used. Dict-style ``names`` and missing names fall back to ``{key}_{i}`` indices.
+    """
+    feature = dataset.features[key]
+    dim = feature["shape"][-1]
+
+    names = feature.get("names")
+    if isinstance(names, list) and len(names) == dim:
+        return [str(name) for name in names]
+
+    return [f"{key}_{d}" for d in range(dim)]
+
+
 def check_chw_float32(frame: torch.Tensor) -> None:
     """
     Check if a frame is a channel-first, float32 tensor.
@@ -93,6 +108,31 @@ def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
     return hwc_uint8_numpy
 
 
+def build_blueprint_from_dataset(dataset: LeRobotDataset):
+    """Build a Rerun blueprint laying out camera images and time series for the given dataset.
+
+    Camera images and scalar signals (action, state, reward, done, success) are arranged in a grid.
+    The per-dimension series names for ``action`` and ``state`` are applied directly
+    via blueprint overrides.
+    """
+    import rerun as rr
+    import rerun.blueprint as rrb
+
+    views = [rrb.Spatial2DView(origin=key, name=key) for key in dataset.meta.camera_keys]
+
+    # Style multi-dimensional signals (action, state) with per-dimension names.
+    for origin, key in ((ACTION, ACTION), ("state", OBS_STATE)):
+        if key in dataset.features:
+            names = get_feature_names(dataset, key)
+            styling = rr.SeriesLines(names=names)
+            views.append(rrb.TimeSeriesView(origin=origin, name=origin, overrides={origin: styling}))
+    for key in (DONE, REWARD, "next.success"):
+        if key in dataset.features:
+            views.append(rrb.TimeSeriesView(origin=key, name=key))
+
+    return rrb.Blueprint(rrb.Grid(*views))
+
+
 def to_hwc_uint16_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
     check_chw_float32(chw_float32_torch)
     hwc_uint16_numpy = chw_float32_torch.round().type(torch.uint16).permute(1, 2, 0).numpy()
@@ -137,7 +177,8 @@ def visualize_dataset(
     import rerun as rr
 
     spawn_local_viewer = mode == "local" and not save
-    rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer)
+    blueprint = build_blueprint_from_dataset(dataset)
+    rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer, default_blueprint=blueprint)
 
     # Manually call python garbage collector after `rr.init` to avoid hanging in a blocking flush
     # when iterating on a dataloader with `num_workers` > 0
@@ -163,12 +204,13 @@ def visualize_dataset(
     for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
         if first_index is None:
             first_index = batch["index"][0].item()
+
         # iterate over the batch
         for i in range(len(batch["index"])):
             rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index)
             rr.set_time("timestamp", timestamp=batch["timestamp"][i].item())
 
-            # display each camera image
+            # display each camera image (or depth map)
             for key in dataset.meta.camera_keys:
                 if key in dataset.meta.depth_keys:
                     depth = to_hwc_uint16_numpy(batch[key][i])
@@ -183,15 +225,13 @@ def visualize_dataset(
                     img_entity = rr.Image(img).compress() if display_compressed_images else rr.Image(img)
                     rr.log(key, entity=img_entity)
 
-            # display each dimension of action space (e.g. actuators command)
+            # display the action space (e.g. actuators command)
             if ACTION in batch:
-                for dim_idx, val in enumerate(batch[ACTION][i]):
-                    rr.log(f"{ACTION}/{dim_idx}", rr.Scalars(val.item()))
+                rr.log(ACTION, rr.Scalars(batch[ACTION][i].numpy()))
 
-            # display each dimension of observed state space (e.g. agent position in joint space)
+            # display the observed state space (e.g. agent position in joint space)
             if OBS_STATE in batch:
-                for dim_idx, val in enumerate(batch[OBS_STATE][i]):
-                    rr.log(f"state/{dim_idx}", rr.Scalars(val.item()))
+                rr.log("state", rr.Scalars(batch[OBS_STATE][i].numpy()))
 
             if DONE in batch:
                 rr.log(DONE, rr.Scalars(batch[DONE][i].item()))
@@ -202,9 +242,8 @@ def visualize_dataset(
             if "next.success" in batch:
                 rr.log("next.success", rr.Scalars(batch["next.success"][i].item()))
 
+    # save .rrd locally
     if mode == "local" and save:
-        # save .rrd locally
-        output_dir = Path(output_dir)
         output_dir.mkdir(parents=True, exist_ok=True)
         repo_id_str = repo_id.replace("/", "_")
         rrd_path = output_dir / f"{repo_id_str}_episode_{episode_index}.rrd"
@@ -212,7 +251,7 @@ def visualize_dataset(
         return rrd_path
 
     elif mode == "distant":
-        # stop the process from exiting since it is serving the websocket connection
+        # Keep the process alive while it serves the gRPC/web connection.
         try:
             while True:
                 time.sleep(1)
@@ -327,12 +366,14 @@ def main():
         )
         logging.warning("Setting grpc_port to ws_port value.")
         kwargs["grpc_port"] = kwargs.pop("ws_port")
+    else:
+        kwargs.pop("ws_port")  # Always remove ws_port from kwargs
 
     init_logging()
     logging.info("Loading dataset")
     dataset = LeRobotDataset(repo_id, episodes=[args.episode_index], root=root, tolerance_s=tolerance_s)
 
-    visualize_dataset(dataset, **vars(args))
+    visualize_dataset(dataset, **kwargs)
 
 
 if __name__ == "__main__":

src/lerobot/scripts/lerobot_train.py

@@ -20,6 +20,7 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
+import sys
 import time
 from contextlib import nullcontext
 from pprint import pformat
@@ -41,15 +42,17 @@ from lerobot.common.train_utils import (
     load_training_batch_size,
     load_training_num_processes,
     load_training_state,
+    push_checkpoint_to_hub,
     save_checkpoint,
     update_last_checkpoint,
 )
 from lerobot.common.wandb_utils import WandBLogger
-from lerobot.configs import parser
+from lerobot.configs import JobConfig, parser
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state
 from lerobot.datasets.factory import make_train_eval_datasets
 from lerobot.envs import close_envs, make_env, make_env_pre_post_processors
+from lerobot.jobs import submit_to_hf
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
 from lerobot.rewards import make_reward_pre_post_processors
@@ -188,6 +191,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         cfg: A `TrainPipelineConfig` object containing all training configurations.
         accelerator: Optional Accelerator instance. If None, one will be created automatically.
     """
+    if cfg.job.is_remote:
+        return submit_to_hf(cfg)
+
     from lerobot.utils.import_utils import require_package
 
     require_package("accelerate", extra="training")
@@ -655,6 +661,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     optim_state_dict=optim_state_dict,
                 )
                 update_last_checkpoint(checkpoint_dir)
+                if cfg.save_checkpoint_to_hub:
+                    push_checkpoint_to_hub(
+                        checkpoint_dir,
+                        cfg.policy.repo_id,
+                        private=cfg.policy.private,
+                    )
                 if wandb_logger:
                     wandb_logger.log_policy(checkpoint_dir)
 
@@ -735,8 +747,25 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     accelerator.end_training()
 
 
+def _remote_target_in_argv() -> bool:
+    """True when the CLI requests a remote HF Jobs run (--job.target=<non-local>)."""
+    target = None
+    args = sys.argv[1:]
+    for i, tok in enumerate(args):
+        if tok == "--job.target" and i + 1 < len(args):
+            target = args[i + 1]
+        elif tok.startswith("--job.target="):
+            target = tok.split("=", 1)[1]
+    return JobConfig.is_remote_target(target)
+
+
 def main():
     register_third_party_plugins()
+    if _remote_target_in_argv():
+        # The policy device is resolved on the remote pod, not here, so silence the
+        # client-side "Device '...' is not available" warning PreTrainedConfig emits
+        # while parsing the config (it fires before train() can dispatch remotely).
+        logging.getLogger("lerobot.configs.policies").setLevel(logging.ERROR)
     train()
 
 

src/lerobot/utils/hub.py

@@ -20,9 +20,33 @@ from typing import Any, TypeVar
 from huggingface_hub import HfApi
 from huggingface_hub.utils import validate_hf_hub_args
 
+from .constants import CHECKPOINTS_DIR
+
 T = TypeVar("T", bound="HubMixin")
 
 
+def find_latest_hub_checkpoint(
+    repo_id: str,
+    *,
+    token: str | bool | None = None,
+    revision: str | None = None,
+) -> str | None:
+    """Repo-relative path of the most recent checkpoint in a training repo.
+
+    Training runs push checkpoints to ``checkpoints/<step>/`` (see
+    ``push_checkpoint_to_hub``). This lists those step dirs and returns
+    ``checkpoints/<highest-step>``, or ``None`` if the repo has no checkpoints.
+    """
+    files = HfApi().list_repo_files(repo_id=repo_id, repo_type="model", revision=revision, token=token)
+    prefix = f"{CHECKPOINTS_DIR}/"
+    steps = {
+        name for f in files if f.startswith(prefix) and (name := f[len(prefix) :].split("/", 1)[0]).isdigit()
+    }
+    if not steps:
+        return None
+    return f"{CHECKPOINTS_DIR}/{max(steps, key=int)}"
+
+
 class HubMixin:
     """
     A Mixin containing the functionality to push an object to the hub.

src/lerobot/utils/visualization_utils.py

@@ -38,6 +38,8 @@ def init_rerun(
     require_package("rerun-sdk", extra="viz", import_name="rerun")
     import rerun as rr
 
+    log_rerun_data.blueprint = None  # Reset blueprint cache for new session
+
     batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000")
     os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size
     rr.init(session_name)
@@ -63,6 +65,41 @@ def _is_scalar(x):
     )
 
 
+def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]):
+    """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views.
+
+    Camera images, observation and action scalars are arranged in a grid.
+    """
+
+    # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
+    import rerun.blueprint as rrb
+
+    views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)]
+
+    if observation_paths:
+        views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths)))
+    if action_paths:
+        views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths)))
+
+    return rrb.Blueprint(rrb.Grid(*views))
+
+
+def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None:
+    """Build and send the blueprint once, from the first observation and action data."""
+    if getattr(log_rerun_data, "blueprint", None) is not None:
+        return
+
+    if not (observation_paths or action_paths or image_paths):
+        return
+
+    # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
+    import rerun as rr
+
+    blueprint = _build_blueprint(observation_paths, action_paths, image_paths)
+    log_rerun_data.blueprint = blueprint
+    rr.send_blueprint(blueprint)
+
+
 def log_rerun_data(
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
@@ -76,11 +113,15 @@ def log_rerun_data(
     - Scalars values (floats, ints) are logged as `rr.Scalars`.
     - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed
       from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`.
-    - 1D NumPy arrays are logged as a series of individual scalars, with each element indexed.
-    - Other multi-dimensional arrays are flattened and logged as individual scalars.
+    - 1D NumPy arrays are logged as a single `rr.Scalars` batch under one entity path, so that every
+      dimension shares the same view instead of being split across one view per element.
+    - Multi-dimensional **action** arrays are flattened and logged as a single `rr.Scalars` batch.
 
     Keys are automatically namespaced with "observation." or "action." if not already present.
 
+    On the first call, a blueprint is built and sent so observation and action scalars get separate
+    time-series views and each image gets its own spatial view.
+
     Args:
         observation: An optional dictionary containing observation data to log.
         action: An optional dictionary containing action data to log.
@@ -90,6 +131,10 @@ def log_rerun_data(
     require_package("rerun-sdk", extra="viz", import_name="rerun")
     import rerun as rr
 
+    observation_paths: set[str] = set()
+    action_paths: set[str] = set()
+    image_paths: set[str] = set()
+
     if observation:
         for k, v in observation.items():
             if v is None:
@@ -98,20 +143,22 @@ def log_rerun_data(
 
             if _is_scalar(v):
                 rr.log(key, rr.Scalars(float(v)))
+                observation_paths.add(key)
             elif isinstance(v, np.ndarray):
                 arr = v
                 # Convert CHW -> HWC when needed
                 if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
                     arr = np.transpose(arr, (1, 2, 0))
                 if arr.ndim == 1:
-                    for i, vi in enumerate(arr):
-                        rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
+                    rr.log(key, rr.Scalars(arr.astype(float)))
+                    observation_paths.add(key)
                 else:
                     if arr.shape[-1] == 1:
                         img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis)
                     else:
                         img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
                     rr.log(key, entity=img_entity, static=True)
+                    image_paths.add(key)
 
     if action:
         for k, v in action.items():
@@ -121,12 +168,10 @@ def log_rerun_data(
 
             if _is_scalar(v):
                 rr.log(key, rr.Scalars(float(v)))
+                action_paths.add(key)
             elif isinstance(v, np.ndarray):
-                if v.ndim == 1:
-                    for i, vi in enumerate(v):
-                        rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
-                else:
-                    # Fall back to flattening higher-dimensional arrays
-                    flat = v.flatten()
-                    for i, vi in enumerate(flat):
-                        rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
+                # Flatten any (incl. higher-dimensional) array into a single batched Scalars
+                rr.log(key, rr.Scalars(v.reshape(-1).astype(float)))
+                action_paths.add(key)
+
+    _ensure_blueprint(observation_paths, action_paths, image_paths)

tests/configs/test_resume_from_hub.py

@@ -0,0 +1,68 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+import lerobot.configs.train as tc
+from lerobot.configs.train import TrainPipelineConfig
+
+
+class _FakeHTTPError(tc.HfHubHTTPError):
+    """HfHubHTTPError that can be raised without a real HTTP response object."""
+
+    def __init__(self):
+        pass
+
+
+def test_from_pretrained_falls_back_to_latest_checkpoint_config(tmp_path, monkeypatch):
+    """A Hub repo with no root train_config.json (an interrupted run that only pushed
+    checkpoints/) resolves via the latest checkpoint's config."""
+    # A real train_config.json written by save_pretrained, to be returned by the fallback.
+    parsed = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"])
+    cfg_file = tmp_path / "train_config.json"
+    parsed._save_pretrained(tmp_path)
+    assert cfg_file.is_file()
+
+    calls = []
+
+    def fake_hf_hub_download(filename=None, **kwargs):
+        calls.append(filename)
+        if filename == "train_config.json":
+            raise _FakeHTTPError()  # no root config
+        if filename == "checkpoints/000010/pretrained_model/train_config.json":
+            return str(cfg_file)
+        raise AssertionError(f"unexpected filename {filename}")
+
+    monkeypatch.setattr(tc, "hf_hub_download", fake_hf_hub_download)
+    monkeypatch.setattr(
+        tc, "find_latest_hub_checkpoint", lambda repo_id, token=None, revision=None: "checkpoints/000010"
+    )
+
+    loaded = TrainPipelineConfig.from_pretrained("user/interrupted-run")
+    assert loaded.dataset.repo_id == "u/d"
+    # Tried the root config first, then fell back to the latest checkpoint's config.
+    assert calls == ["train_config.json", "checkpoints/000010/pretrained_model/train_config.json"]
+
+
+def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypatch):
+    """No root config AND no checkpoints → a clear FileNotFoundError, not the raw HTTP error."""
+
+    def fake_hf_hub_download(filename=None, **kwargs):
+        raise _FakeHTTPError()
+
+    monkeypatch.setattr(tc, "hf_hub_download", fake_hf_hub_download)
+    monkeypatch.setattr(tc, "find_latest_hub_checkpoint", lambda repo_id, token=None, revision=None: None)
+
+    with pytest.raises(FileNotFoundError, match="train_config.json not found"):
+        TrainPipelineConfig.from_pretrained("user/empty-repo")

tests/jobs/conftest.py

@@ -0,0 +1,17 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Importing concrete policy configs registers their draccus `--policy.type`
+# choices (e.g. "act") so tests can parse them.
+from lerobot.policies.act.configuration_act import ACTConfig  # noqa: F401

tests/jobs/test_dataset.py

@@ -0,0 +1,66 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock
+
+import pytest
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.jobs.dataset import ensure_dataset_available
+
+
+def _api_with_dataset(exists: bool):
+    api = MagicMock()
+    api.repo_exists.return_value = exists
+    return api
+
+
+def _make_local_cache(tmp_path, repo_id: str) -> None:
+    """Create the minimal local-cache layout that ensure_dataset_available checks."""
+    info = tmp_path / repo_id / "meta" / "info.json"
+    info.parent.mkdir(parents=True)
+    info.write_text("{}")
+
+
+# Branch 1: dataset already on Hub → no push, no error (pod downloads by repo_id).
+def test_dataset_already_on_hub_is_noop():
+    api = _api_with_dataset(True)
+    assert ensure_dataset_available("user/ds", api=api) is None
+    api.repo_exists.assert_called_once_with("user/ds", repo_type="dataset")
+
+
+# Branch 2: not on Hub but present locally → always push privately.
+def test_dataset_local_only_uploads_privately(tmp_path, monkeypatch):
+    monkeypatch.setattr("lerobot.jobs.dataset.HF_LEROBOT_HOME", tmp_path)
+    _make_local_cache(tmp_path, "user/ds")
+
+    api = _api_with_dataset(False)
+    mock_ds_cls = MagicMock()
+    monkeypatch.setattr("lerobot.jobs.dataset.LeRobotDataset", mock_ds_cls)
+
+    assert ensure_dataset_available("user/ds", api=api, tags=["lerobot", "lelab"]) is None
+
+    mock_ds_cls.assert_called_once_with("user/ds")
+    mock_ds_cls.return_value.push_to_hub.assert_called_once_with(private=True, tags=["lerobot", "lelab"])
+
+
+# Branch 3: not on Hub, NOT in local cache → RuntimeError.
+def test_dataset_neither_on_hub_nor_local_raises(tmp_path, monkeypatch):
+    monkeypatch.setattr("lerobot.jobs.dataset.HF_LEROBOT_HOME", tmp_path)
+    # tmp_path is empty — no local cache.
+
+    api = _api_with_dataset(False)
+    with pytest.raises(RuntimeError, match="not in the local cache"):
+        ensure_dataset_available("user/ds", api=api)

tests/jobs/test_hf.py

@@ -0,0 +1,493 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime as dt
+import json
+import threading
+from types import SimpleNamespace
+
+import draccus
+import httpx
+import pytest
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.jobs.hf import (
+    _pod_forwarded_args,
+    _poll_until_done,
+    build_remote_config_file,
+    build_repo_id,
+    resolve_job_tags,
+    resolve_wandb_api_key,
+    submit_to_hf,
+)
+
+
+def test_resolve_job_tags_always_includes_lerobot_and_dedups():
+    assert resolve_job_tags(None) == ["lerobot"]
+    assert resolve_job_tags([]) == ["lerobot"]
+    assert resolve_job_tags(["lelab"]) == ["lerobot", "lelab"]
+    # lerobot isn't duplicated if passed explicitly; order is stable.
+    assert resolve_job_tags(["lelab", "lerobot", "lelab"]) == ["lerobot", "lelab"]
+
+
+def _fake_inspect(stage_value, *, as_enum=True):
+    # huggingface_hub returns `stage` as an enum (with `.value`) in some versions and a plain str in others.
+    stage = SimpleNamespace(value=stage_value) if as_enum else stage_value
+    return lambda job_id: SimpleNamespace(status=SimpleNamespace(stage=stage))
+
+
+@pytest.mark.parametrize("as_enum", [True, False], ids=["enum_stage", "str_stage"])
+def test_poll_until_done_returns_terminal_stage(monkeypatch, as_enum):
+    monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _fake_inspect("COMPLETED", as_enum=as_enum))
+    done = threading.Event()
+    assert _poll_until_done("j", done, poll_interval=0.01) == "COMPLETED"
+    assert done.is_set()
+
+
+def test_poll_until_done_exits_when_done_already_set(monkeypatch):
+    # Non-terminal forever; with done pre-set the loop must not block and returns None.
+    monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _fake_inspect("RUNNING"))
+    done = threading.Event()
+    done.set()
+    assert _poll_until_done("j", done, poll_interval=0.01) is None
+
+
+def test_poll_until_done_gives_up_after_repeated_network_failures(monkeypatch):
+    monkeypatch.setattr(
+        "lerobot.jobs.hf.inspect_job", lambda job_id: (_ for _ in ()).throw(httpx.ConnectError("boom"))
+    )
+    done = threading.Event()
+    result = _poll_until_done("j", done, poll_interval=0.001, max_failures=3)
+    assert result is None
+    assert done.is_set()
+
+
+def test_poll_until_done_propagates_programming_errors(monkeypatch):
+    """A bug (e.g. TypeError) must surface, not be silently retried as a transient failure."""
+    monkeypatch.setattr("lerobot.jobs.hf.inspect_job", lambda job_id: (_ for _ in ()).throw(TypeError("bug")))
+    done = threading.Event()
+    with pytest.raises(TypeError):
+        _poll_until_done("j", done, poll_interval=0.001, max_failures=3)
+
+
+def test_resolve_wandb_key_from_env(monkeypatch):
+    monkeypatch.setenv("WANDB_API_KEY", "abc123")
+    assert resolve_wandb_api_key() == "abc123"
+
+
+def test_resolve_wandb_key_missing(monkeypatch, tmp_path):
+    monkeypatch.delenv("WANDB_API_KEY", raising=False)
+    monkeypatch.setenv("HOME", str(tmp_path))  # no ~/.netrc here
+    monkeypatch.setattr("netrc.netrc", lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError()))
+    assert resolve_wandb_api_key() is None
+
+
+def test_resolve_wandb_key_from_netrc(monkeypatch):
+    # No env var → fall back to the wandb credentials in ~/.netrc.
+    monkeypatch.delenv("WANDB_API_KEY", raising=False)
+
+    class _FakeNetrc:
+        def authenticators(self, host):
+            assert host == "api.wandb.ai"
+            return ("login", "account", "netrc-secret")
+
+    monkeypatch.setattr("netrc.netrc", lambda *a, **k: _FakeNetrc())
+    assert resolve_wandb_api_key() == "netrc-secret"
+
+
+def test_resolve_wandb_key_netrc_without_wandb_entry(monkeypatch):
+    # ~/.netrc exists but has no api.wandb.ai entry → None.
+    monkeypatch.delenv("WANDB_API_KEY", raising=False)
+
+    class _FakeNetrc:
+        def authenticators(self, host):
+            return None
+
+    monkeypatch.setattr("netrc.netrc", lambda *a, **k: _FakeNetrc())
+    assert resolve_wandb_api_key() is None
+
+
+def test_build_repo_id_sanitizes_and_timestamps():
+    now = dt.datetime(2026, 6, 19, 10, 22, 3)
+    assert build_repo_id("alice", "act", now) == "alice/act_2026-06-19_10-22-03"
+    # Runs of illegal characters collapse to a single dash; edges are trimmed.
+    assert build_repo_id("alice", "my cool/run!!", now) == "alice/my-cool-run_2026-06-19_10-22-03"
+    # A name with nothing usable falls back to "train".
+    assert build_repo_id("alice", "///", now) == "alice/train_2026-06-19_10-22-03"
+
+
+def test_pod_forwarded_args_drops_host_only_flags():
+    """User overrides are replayed on the pod, minus flags that only make sense on the submitter.
+
+    `--dataset.root` is a host-local path the pod can't read, so it must be dropped in both the
+    `--name=value` and `--name value` forms; unrelated overrides are forwarded untouched.
+    """
+    argv = [
+        "--config_path=u/d",
+        "--dataset.root=/local/data",
+        "--dataset.root",
+        "/other/local/data",
+        "--policy.repo_id=u/keep",
+        "--steps=10",
+        "--job.target=a10g-small",
+    ]
+    forwarded = _pod_forwarded_args(
+        argv,
+        drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"),
+        drop_prefixes=("--job.",),
+    )
+    assert forwarded == ["--steps=10"]
+
+
+def _minimal_cfg():
+    return draccus.parse(
+        TrainPipelineConfig,
+        args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
+    )
+
+
+def test_validate_skips_repo_id_check_for_remote():
+    """Remote runs auto-assign repo_id in submit_to_hf, so validate() must not demand it up front."""
+    cfg = _minimal_cfg()  # remote target, push_to_hub default True, no explicit repo_id
+    assert cfg.policy.repo_id is None
+    cfg.validate()  # must not raise
+
+
+def test_validate_requires_repo_id_for_local_push():
+    """Local runs that push to the Hub still need an explicit repo_id."""
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=["--dataset.repo_id", "u/d", "--policy.type", "act"],
+    )
+    with pytest.raises(ValueError, match="repo_id"):
+        cfg.validate()
+
+
+def test_build_remote_config_applies_overrides(tmp_path):
+    cfg = _minimal_cfg()
+    dest = tmp_path / "train_config.json"
+    out = build_remote_config_file(cfg, "u/run", dest)
+    assert out == dest
+    data = json.loads(dest.read_text())
+    # `job` is client-only orchestration and must be stripped for the pod.
+    assert "job" not in data
+    # save_checkpoint_to_hub defaults off → omitted so older images accept the config.
+    assert "save_checkpoint_to_hub" not in data
+    assert data["policy"]["push_to_hub"] is True
+    assert data["policy"]["repo_id"] == "u/run"
+    assert data["policy"]["device"] is None  # pod auto-detects its GPU
+    assert data["dataset"]["root"] is None  # pod resolves the dataset by repo_id
+    # the caller's cfg must be left untouched (function works on a deep copy)
+    assert cfg.job.target == "a10g-small"
+    assert cfg.save_checkpoint_to_hub is False
+
+
+def test_build_remote_config_includes_checkpoint_flag_when_enabled(tmp_path):
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=[
+            "--dataset.repo_id",
+            "u/d",
+            "--policy.type",
+            "act",
+            "--job.target",
+            "a10g-small",
+            "--save_checkpoint_to_hub",
+            "true",
+        ],
+    )
+    dest = tmp_path / "train_config.json"
+    build_remote_config_file(cfg, "u/run", dest)
+    data = json.loads(dest.read_text())
+    # explicitly enabled → kept in the config (requires a matching trainer image).
+    assert data["save_checkpoint_to_hub"] is True
+    assert "job" not in data
+
+
+def test_build_remote_config_merges_tags_into_policy(tmp_path):
+    cfg = _minimal_cfg()
+    dest = tmp_path / "train_config.json"
+    build_remote_config_file(cfg, "u/run", dest, tags=["lerobot", "lelab"])
+    data = json.loads(dest.read_text())
+    # tags propagate to the model the pod pushes.
+    assert data["policy"]["tags"] == ["lerobot", "lelab"]
+
+
+def test_build_remote_config_merges_tags_without_duplicating(tmp_path):
+    cfg = _minimal_cfg()
+    cfg.policy.tags = ["existing", "lerobot"]
+    dest = tmp_path / "train_config.json"
+    build_remote_config_file(cfg, "u/run", dest, tags=["lerobot", "lelab"])
+    data = json.loads(dest.read_text())
+    # pre-existing policy tags are kept; only genuinely-new tags are appended (no dup "lerobot").
+    assert data["policy"]["tags"] == ["existing", "lerobot", "lelab"]
+
+
+def test_submit_requires_login(monkeypatch):
+    monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: None)
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
+    )
+    with pytest.raises(RuntimeError, match="hf auth login"):
+        submit_to_hf(cfg)
+
+
+def test_submit_passes_validation_and_submits(monkeypatch):
+    """A type-based policy with no explicit repo_id is auto-assigned one and submitted."""
+    from unittest.mock import MagicMock
+
+    # Patch get_token
+    monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
+
+    # Patch HfApi so whoami returns alice
+    class FakeHfApi:
+        def __init__(self, token=None):
+            pass
+
+        def whoami(self, token=None):
+            return {"name": "alice"}
+
+    monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
+
+    # ensure_dataset_available returns None; patch it out so no Hub access happens
+    # (hf.py imports it at module level, so patch it on lerobot.jobs.hf).
+    monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None)
+
+    # Patch _stage_config_on_hub to skip network
+    monkeypatch.setattr(
+        "lerobot.jobs.hf._stage_config_on_hub",
+        lambda cfg, repo_id, token, tags=None: repo_id,
+    )
+
+    # Patch run_job to return a fake job
+    fake_job = MagicMock()
+    fake_job.id = "job-123"
+    run_job_calls = []
+
+    def fake_run_job(**kwargs):
+        run_job_calls.append(kwargs)
+        return fake_job
+
+    monkeypatch.setattr("lerobot.jobs.hf.run_job", fake_run_job)
+
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=[
+            "--dataset.repo_id",
+            "u/d",
+            "--policy.type",
+            "act",
+            "--job.target",
+            "a10g-small",
+            "--job.detach",
+            "true",
+        ],
+    )
+
+    # Must NOT raise (pre-fix this raised ValueError about missing repo_id)
+    submit_to_hf(cfg)
+
+    assert len(run_job_calls) == 1, "run_job should have been called exactly once"
+    assert cfg.policy.repo_id is not None
+    assert cfg.policy.repo_id.startswith("alice/")
+    call = run_job_calls[0]
+    # The pod runs `lerobot-train --config_path=<staged repo>` on the requested flavor/image.
+    assert call["command"][0] == "lerobot-train"
+    assert call["command"][1].startswith("--config_path=")
+    assert call["flavor"] == "a10g-small"
+    assert call["image"] == "huggingface/lerobot-gpu:latest"
+    # The Hub token is forwarded so the pod can pull the (possibly private) dataset.
+    assert call["secrets"]["HF_TOKEN"] == "tok"
+    # Every job carries the lerobot tag as a queryable label.
+    assert call["labels"].get("lerobot") == "true"
+
+
+def test_submit_rejects_reward_model_training(monkeypatch):
+    """Remote training only supports policies; reward-model runs fail fast with a clear error."""
+    monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
+
+    class FakeHfApi:
+        def __init__(self, token=None):
+            pass
+
+        def whoami(self, token=None):
+            return {"name": "alice"}
+
+    monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
+
+    cfg = _minimal_cfg()
+    cfg.reward_model = SimpleNamespace(type="reward")  # marks this as reward-model training
+    monkeypatch.setattr(cfg, "validate", lambda: None)  # skip pretrained-path resolution
+
+    with pytest.raises(ValueError, match="reward model"):
+        submit_to_hf(cfg)
+
+
+@pytest.mark.timeout(15)
+def test_submit_returns_when_job_completes(monkeypatch):
+    """Non-detach path must RETURN (not hang) once the job reaches a terminal stage."""
+    from types import SimpleNamespace
+
+    monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
+
+    class FakeHfApi:
+        def __init__(self, token=None):
+            pass
+
+        def whoami(self, token=None):
+            return {"name": "alice"}
+
+    monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
+    monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None)
+    monkeypatch.setattr(
+        "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id
+    )
+    monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x"))
+    # Job is already COMPLETED on the first poll.
+    monkeypatch.setattr(
+        "lerobot.jobs.hf.inspect_job",
+        lambda job_id: SimpleNamespace(
+            status=SimpleNamespace(stage=SimpleNamespace(value="COMPLETED"), message=None)
+        ),
+    )
+    # Log stream ends immediately.
+    monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
+
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
+    )
+    # Runs in the pytest main thread (signal handler install requires it); the
+    # @timeout marker fails the test instead of hanging if it regresses.
+    submit_to_hf(cfg)
+
+
+@pytest.mark.timeout(15)
+def test_submit_returns_on_model_pushed_marker(monkeypatch):
+    """Finish when the model-pushed log appears, even if the job stage never flips."""
+    from types import SimpleNamespace
+
+    monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
+
+    class FakeHfApi:
+        def __init__(self, token=None):
+            pass
+
+        def whoami(self, token=None):
+            return {"name": "alice"}
+
+    monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
+    monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None)
+    monkeypatch.setattr(
+        "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id
+    )
+    monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x"))
+    # Job stays RUNNING forever — only the log marker can end the command.
+    monkeypatch.setattr(
+        "lerobot.jobs.hf.inspect_job",
+        lambda job_id: SimpleNamespace(
+            status=SimpleNamespace(stage=SimpleNamespace(value="RUNNING"), message=None)
+        ),
+    )
+    pushed_line = "INFO Model pushed to https://huggingface.co/alice/myrun"
+    monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter([pushed_line]))
+
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=[
+            "--dataset.repo_id",
+            "u/d",
+            "--policy.type",
+            "act",
+            "--policy.repo_id",
+            "alice/myrun",
+            "--job.target",
+            "a10g-small",
+        ],
+    )
+    # Must return via the model-pushed marker despite the perpetual RUNNING stage.
+    submit_to_hf(cfg)
+
+
+def test_submit_raises_when_wandb_enabled_without_key(monkeypatch):
+    """wandb.enable with no key reachable anywhere fails fast, before submitting."""
+
+    monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
+
+    class FakeHfApi:
+        def __init__(self, token=None):
+            pass
+
+        def whoami(self, token=None):
+            return {"name": "alice"}
+
+    monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
+    monkeypatch.setattr("lerobot.jobs.hf.resolve_wandb_api_key", lambda: None)
+
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=[
+            "--dataset.repo_id",
+            "u/d",
+            "--policy.type",
+            "act",
+            "--job.target",
+            "a10g-small",
+            "--wandb.enable",
+            "true",
+        ],
+    )
+    with pytest.raises(ValueError, match="WANDB_API_KEY"):
+        submit_to_hf(cfg)
+
+
+@pytest.mark.timeout(15)
+def test_submit_raises_when_job_ends_in_error(monkeypatch):
+    """A terminal non-COMPLETED stage with no model-pushed marker must raise with the status."""
+    from types import SimpleNamespace
+
+    monkeypatch.setattr("lerobot.jobs.hf.get_token", lambda: "tok")
+
+    class FakeHfApi:
+        def __init__(self, token=None):
+            pass
+
+        def whoami(self, token=None):
+            return {"name": "alice"}
+
+    monkeypatch.setattr("lerobot.jobs.hf.HfApi", FakeHfApi)
+    monkeypatch.setattr("lerobot.jobs.hf.ensure_dataset_available", lambda *a, **kw: None)
+    monkeypatch.setattr(
+        "lerobot.jobs.hf._stage_config_on_hub", lambda cfg, repo_id, token, tags=None: repo_id
+    )
+    monkeypatch.setattr("lerobot.jobs.hf.run_job", lambda **kw: SimpleNamespace(id="job-1", url="http://x"))
+    # Job fails: a terminal ERROR stage carrying the platform's status message.
+    monkeypatch.setattr(
+        "lerobot.jobs.hf.inspect_job",
+        lambda job_id: SimpleNamespace(
+            status=SimpleNamespace(stage=SimpleNamespace(value="ERROR"), message="Job timeout")
+        ),
+    )
+    # Logs end without the model-pushed marker.
+    monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
+
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
+    )
+    with pytest.raises(RuntimeError, match=r"stage=ERROR \(Job timeout\)"):
+        submit_to_hf(cfg)

tests/jobs/test_job_config.py

@@ -0,0 +1,64 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import draccus
+import pytest
+
+from lerobot.configs import JobConfig
+from lerobot.configs.train import TrainPipelineConfig
+
+
+def test_jobconfig_defaults_are_local():
+    cfg = JobConfig()
+    assert cfg.target is None
+    assert cfg.is_remote is False
+    assert cfg.image == "huggingface/lerobot-gpu:latest"
+    assert cfg.timeout == "2d"
+    assert cfg.detach is False
+
+
+def test_jobconfig_local_string_is_not_remote():
+    assert JobConfig(target="local").is_remote is False
+
+
+def test_jobconfig_flavor_is_remote():
+    assert JobConfig(target="a10g-small").is_remote is True
+
+
+def test_train_config_parses_job_target():
+    parsed = draccus.parse(
+        TrainPipelineConfig,
+        args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
+    )
+    assert parsed.job.target == "a10g-small"
+    assert parsed.job.is_remote is True
+    assert parsed.save_checkpoint_to_hub is False
+
+
+def test_save_checkpoint_to_hub_requires_repo_id():
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=[
+            "--dataset.repo_id",
+            "u/d",
+            "--policy.type",
+            "act",
+            "--policy.push_to_hub",
+            "false",
+            "--save_checkpoint_to_hub",
+            "true",
+        ],
+    )
+    with pytest.raises(ValueError, match="requires --policy.repo_id"):
+        cfg.validate()

tests/policies/molmoact2/test_molmoact2.py

@@ -44,10 +44,12 @@ from lerobot.policies.molmoact2.modeling_molmoact2 import (
     _combine_rollout_seeds,
 )
 from lerobot.policies.molmoact2.processor_molmoact2 import (
+    MolmoAct2ActionFrameTransformStep,
     MolmoAct2ClampNormalizedProcessorStep,
     MolmoAct2MaskedNormalizerProcessorStep,
     MolmoAct2MaskedUnnormalizerProcessorStep,
     MolmoAct2PackInputsProcessorStep,
+    MolmoAct2StateFrameTransformStep,
     _add_gripper_masks_to_stats,
     _build_discrete_state_string,
     _normalize_question_text,
@@ -926,6 +928,39 @@ def test_question_normalization_matches_release_prompt_style():
     )
 
 
+def test_joint_frame_transform_round_trip():
+    signs = [1.0, -1.0, 1.0, 1.0, 1.0, 1.0]
+    offsets = [0.0, 90.0, 90.0, 0.0, 0.0, 0.0]
+    original_state = torch.tensor([[10.0, -90.0, -120.0, 30.0, 0.0, -45.0]])
+
+    state_step = MolmoAct2StateFrameTransformStep(joint_signs=signs, joint_offsets=offsets)
+    action_step = MolmoAct2ActionFrameTransformStep(joint_signs=signs, joint_offsets=offsets)
+
+    transition = {
+        TransitionKey.OBSERVATION: {OBS_STATE: original_state.clone()},
+    }
+    transformed = state_step(transition)
+    model_state = transformed[TransitionKey.OBSERVATION][OBS_STATE]
+
+    action_transition = {TransitionKey.ACTION: model_state.clone()}
+    recovered = action_step(action_transition)
+    recovered_state = recovered[TransitionKey.ACTION]
+
+    assert torch.allclose(recovered_state, original_state)
+
+
+def test_joint_frame_transform_noop_when_none():
+    state_step = MolmoAct2StateFrameTransformStep(joint_signs=None, joint_offsets=None)
+    action_step = MolmoAct2ActionFrameTransformStep(joint_signs=None, joint_offsets=None)
+    state = torch.tensor([[10.0, -90.0, -120.0]])
+
+    state_transition = {TransitionKey.OBSERVATION: {OBS_STATE: state}}
+    assert state_step(state_transition) is state_transition
+
+    action_transition = {TransitionKey.ACTION: state}
+    assert action_step(action_transition) is action_transition
+
+
 def test_action_padding_marks_only_real_dimensions():
     step = object.__new__(MolmoAct2PackInputsProcessorStep)
     step.max_action_dim = 32

tests/policies/vla_jepa/conftest.py

@@ -8,7 +8,6 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
-from PIL import Image
 from torch import Tensor, nn
 
 from lerobot.configs.types import FeatureType, PolicyFeature
@@ -191,7 +190,7 @@ class _FakeQwenInterface(nn.Module):
 
     def build_inputs(
         self,
-        images: list[list[Image.Image]],
+        images: list[list[Tensor]],
         instructions: list[str],
         action_prompt: str,
         embodied_prompt: str,
@@ -214,12 +213,13 @@ class _FakeQwenInterface(nn.Module):
         }
 
     @staticmethod
-    def tensor_to_pil(image_tensor: Tensor) -> Image.Image:
-        image = image_tensor.detach().cpu()
-        if image.ndim == 3 and image.shape[0] in (1, 3):
-            image = image.permute(1, 2, 0)
-        image = (image.float().clamp(0, 1) * 255).to(torch.uint8).numpy()
-        return Image.fromarray(image)
+    def to_pixel_values(image_tensor: Tensor) -> Tensor:
+        image = image_tensor.detach().float()
+        if image.shape[-3] == 1:
+            repeats = [1] * image.ndim
+            repeats[-3] = 3
+            image = image.repeat(*repeats)
+        return image
 
 
 class _FakeVideoEncoder(nn.Module):
@@ -242,12 +242,14 @@ class _FakeVideoEncoder(nn.Module):
 
 
 class _FakeVideoProcessor:
-    def __call__(self, videos, return_tensors: str) -> dict[str, Tensor]:
+    def __call__(self, videos, return_tensors: str, device=None, **kwargs) -> dict[str, Tensor]:
         assert return_tensors == "pt"
         if isinstance(videos, list):
             pixel_values = torch.stack([torch.as_tensor(v) for v in videos])
         else:
             pixel_values = torch.as_tensor(videos).unsqueeze(0)
+        if device is not None:
+            pixel_values = pixel_values.to(device)
         return {"pixel_values_videos": pixel_values}
 
 

tests/policies/vla_jepa/test_vla_jepa.py

@@ -211,40 +211,42 @@ def test_reset_clears_action_queue(patch_vla_jepa_external_models: None) -> None
 
 
 def test_prepare_model_inputs_training_format(patch_vla_jepa_external_models: None) -> None:
-    from PIL import Image
-
     policy = VLAJEPAPolicy(make_config())
-    examples = policy._prepare_model_inputs(make_train_batch())
+    inputs = policy._prepare_model_inputs(make_train_batch())
 
-    assert len(examples) == BATCH_SIZE
-    for ex in examples:
-        assert set(ex) >= {"image", "video", "lang", "action", "state"}
-        assert len(ex["image"]) == 1 and isinstance(ex["image"][0], Image.Image)
-        assert ex["video"].ndim == 5 and ex["video"].dtype == np.uint8  # [V,T,H,W,C]
-        assert ex["action"].shape == (ACTION_HORIZON, ACTION_DIM)
-        assert ex["state"].shape == (1, STATE_DIM)
+    assert set(inputs) >= {"images", "instructions", "videos", "actions", "state"}
+    # images: per-sample, per-view [C, H, W] float tensors (kept as a list for Qwen messages)
+    assert len(inputs["images"]) == BATCH_SIZE and len(inputs["images"][0]) == 1
+    img = inputs["images"][0][0]
+    assert isinstance(img, torch.Tensor) and img.dtype == torch.float32 and img.ndim == 3
+    assert len(inputs["instructions"]) == BATCH_SIZE
+    # videos: batched [B, V, T, C, H, W] float
+    assert inputs["videos"].ndim == 6 and inputs["videos"].shape[0] == BATCH_SIZE
+    assert inputs["videos"].dtype == torch.float32
+    assert inputs["actions"].shape == (BATCH_SIZE, ACTION_HORIZON, ACTION_DIM)
+    assert inputs["state"].shape == (BATCH_SIZE, 1, STATE_DIM)
 
 
 def test_prepare_model_inputs_inference_omits_action(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
-    for ex in policy._prepare_model_inputs(make_inference_batch()):
-        assert "action" not in ex
-        assert "image" in ex and "video" in ex and "lang" in ex
+    inputs = policy._prepare_model_inputs(make_inference_batch())
+    assert "actions" not in inputs and "action_is_pad" not in inputs
+    assert {"images", "instructions", "state"} <= set(inputs)
 
 
 def test_prepare_model_inputs_missing_task_uses_default(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     del batch["task"]
-    examples = policy._prepare_model_inputs(batch)
-    assert all(isinstance(ex["lang"], str) and len(ex["lang"]) > 0 for ex in examples)
+    instructions = policy._prepare_model_inputs(batch)["instructions"]
+    assert all(isinstance(s, str) and len(s) > 0 for s in instructions)
 
 
 def test_prepare_model_inputs_string_task_broadcast(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     batch["task"] = "open the drawer"
-    assert all(ex["lang"] == "open the drawer" for ex in policy._prepare_model_inputs(batch))
+    assert policy._prepare_model_inputs(batch)["instructions"] == ["open the drawer"] * BATCH_SIZE
 
 
 def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: None) -> None:
@@ -253,7 +255,7 @@ def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: N
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     del batch[OBS_STATE]
-    assert all("state" not in ex for ex in policy._prepare_model_inputs(batch))
+    assert "state" not in policy._prepare_model_inputs(batch)
 
 
 # ---------------------------------------------------------------------------
@@ -446,14 +448,14 @@ def test_postprocessor_applied_after_predict_action_chunk(
     """
     from lerobot.policies.vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
 
-    raw_actions = np.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=np.float32)
+    raw_actions = torch.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=torch.float32)
 
     cfg = make_config()
     cfg.clip_normalized_actions = False
     cfg.binarize_gripper_action = False
     policy = VLAJEPAPolicy(cfg)
     policy.eval()
-    monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.copy())
+    monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.clone())
 
     dataset_stats = _make_dataset_stats()
     _, postprocessor = make_vla_jepa_pre_post_processors(cfg, dataset_stats)
@@ -564,9 +566,9 @@ def test_single_view_is_duplicated_for_world_model(patch_vla_jepa_external_model
     original_processor = policy.model.video_processor
 
     class _CapturingProcessor:
-        def __call__(self, videos: list, return_tensors: str) -> dict:
+        def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict:
             captured_videos.extend(videos)
-            return original_processor(videos=videos, return_tensors=return_tensors)
+            return original_processor(videos=videos, return_tensors=return_tensors, **kwargs)
 
     policy.model.video_processor = _CapturingProcessor()
     policy.forward(_make_multiview_train_batch(num_views=1))
@@ -587,9 +589,9 @@ def test_excess_views_trimmed_for_world_model(patch_vla_jepa_external_models: No
     original_processor = policy.model.video_processor
 
     class _CapturingProcessor:
-        def __call__(self, videos: list, return_tensors: str) -> dict:
+        def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict:
             captured_videos.extend(videos)
-            return original_processor(videos=videos, return_tensors=return_tensors)
+            return original_processor(videos=videos, return_tensors=return_tensors, **kwargs)
 
     policy.model.video_processor = _CapturingProcessor()
     policy.forward(_make_multiview_train_batch(num_views=3))

tests/scripts/test_train_remote_dispatch.py

@@ -0,0 +1,67 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+
+import draccus
+import pytest
+
+# Importing lerobot_train eagerly pulls in lerobot.datasets, which needs the
+# `dataset` extra. The base CI tier runs without it, so skip the whole module there.
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.configs.train import TrainPipelineConfig  # noqa: E402
+from lerobot.policies.act.configuration_act import (
+    ACTConfig,  # noqa: E402, F401  (registers --policy.type act)
+)
+from lerobot.scripts.lerobot_train import _remote_target_in_argv, train  # noqa: E402
+
+
+def _set_argv(monkeypatch, *args):
+    monkeypatch.setattr(sys, "argv", ["lerobot-train", *args])
+
+
+def test_remote_target_detected_space_separated(monkeypatch):
+    _set_argv(monkeypatch, "--policy.type", "act", "--job.target", "a10g-small")
+    assert _remote_target_in_argv() is True
+
+
+def test_remote_target_detected_equals(monkeypatch):
+    _set_argv(monkeypatch, "--job.target=t4-small")
+    assert _remote_target_in_argv() is True
+
+
+def test_local_string_is_not_remote(monkeypatch):
+    _set_argv(monkeypatch, "--job.target", "local")
+    assert _remote_target_in_argv() is False
+
+
+def test_no_target_is_not_remote(monkeypatch):
+    _set_argv(monkeypatch, "--policy.type", "act")
+    assert _remote_target_in_argv() is False
+
+
+def test_train_dispatches_to_submit_when_remote(monkeypatch):
+    """A remote --job.target short-circuits train() to the HF Jobs submitter."""
+    import lerobot.scripts.lerobot_train as train_module
+
+    captured = []
+    monkeypatch.setattr(train_module, "submit_to_hf", lambda cfg: captured.append(cfg) or "submitted")
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
+    )
+    # Returns the submitter's result and never enters the local training path.
+    assert train(cfg) == "submitted"
+    assert captured == [cfg]

tests/utils/test_hub.py

@@ -0,0 +1,54 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import MagicMock
+
+from lerobot.utils.hub import find_latest_hub_checkpoint
+
+
+def _patch_list_files(monkeypatch, files):
+    api = MagicMock()
+    api.list_repo_files.return_value = files
+    # HfApi is imported into lerobot.utils.hub at module load, so patch it there.
+    monkeypatch.setattr("lerobot.utils.hub.HfApi", lambda *a, **k: api)
+    return api
+
+
+def test_find_latest_hub_checkpoint_picks_highest_step(monkeypatch):
+    _patch_list_files(
+        monkeypatch,
+        [
+            "README.md",
+            "checkpoints/000500/pretrained_model/model.safetensors",
+            "checkpoints/000500/training_state/training_step.json",
+            "checkpoints/020000/pretrained_model/model.safetensors",
+            "checkpoints/001000/training_state/training_step.json",
+        ],
+    )
+    # Numeric max, not lexicographic — "020000" beats "001000"/"000500".
+    assert find_latest_hub_checkpoint("u/run") == "checkpoints/020000"
+
+
+def test_find_latest_hub_checkpoint_ignores_non_step_entries(monkeypatch):
+    _patch_list_files(
+        monkeypatch,
+        ["checkpoints/last/pretrained_model/model.safetensors", "config.json"],
+    )
+    # "last" (a symlink target name) is not a numeric step → no resolvable checkpoint.
+    assert find_latest_hub_checkpoint("u/run") is None
+
+
+def test_find_latest_hub_checkpoint_none_when_no_checkpoints(monkeypatch):
+    _patch_list_files(monkeypatch, ["config.json", "model.safetensors"])
+    assert find_latest_hub_checkpoint("u/run") is None

tests/utils/test_train_utils.py

@@ -15,7 +15,9 @@
 # limitations under the License.
 
 from pathlib import Path
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
 
 from lerobot.common.train_utils import (
     get_step_checkpoint_dir,
@@ -24,6 +26,7 @@ from lerobot.common.train_utils import (
     load_training_num_processes,
     load_training_state,
     load_training_step,
+    push_checkpoint_to_hub,
     save_checkpoint,
     save_training_state,
     save_training_step,
@@ -151,3 +154,72 @@ def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler):
     assert loaded_step == 10
     assert loaded_optimizer is optimizer
     assert loaded_scheduler is scheduler
+
+
+def test_push_checkpoint_to_hub_creates_repo_and_uploads(tmp_path, monkeypatch):
+    ckpt = tmp_path / "010000"
+    (ckpt / "pretrained_model").mkdir(parents=True)
+    api = MagicMock()
+    monkeypatch.setattr("lerobot.common.train_utils.HfApi", lambda *a, **k: api)
+    push_checkpoint_to_hub(ckpt, "user/run", private=True)
+    api.create_repo.assert_called_once()
+    assert api.create_repo.call_args.kwargs["private"] is True
+    assert api.create_repo.call_args.kwargs["repo_type"] == "model"
+    api.upload_folder.assert_called_once()
+    kwargs = api.upload_folder.call_args.kwargs
+    assert kwargs["repo_id"] == "user/run"
+    assert kwargs["repo_type"] == "model"
+    assert kwargs["path_in_repo"] == "checkpoints/010000"
+    assert kwargs["folder_path"] == str(ckpt)
+    assert kwargs["commit_message"] == "checkpoint 010000"
+    # A tag named after the checkpoint step is created so the checkpoint can be
+    # recovered with --policy.pretrained_revision instead of a commit sha.
+    api.create_tag.assert_called_once()
+    tag_kwargs = api.create_tag.call_args.kwargs
+    assert tag_kwargs["tag"] == "010000"
+    assert tag_kwargs["revision"] == api.upload_folder.return_value.oid
+    assert tag_kwargs["repo_type"] == "model"
+    assert tag_kwargs["exist_ok"] is True
+
+
+def test_push_checkpoint_to_hub_defaults_to_hub_default_visibility(tmp_path, monkeypatch):
+    ckpt = tmp_path / "010000"
+    (ckpt / "pretrained_model").mkdir(parents=True)
+    api = MagicMock()
+    monkeypatch.setattr("lerobot.common.train_utils.HfApi", lambda *a, **k: api)
+    push_checkpoint_to_hub(ckpt, "user/run")
+    api.create_repo.assert_called_once()
+    assert api.create_repo.call_args.kwargs["private"] is None
+
+
+def test_resolve_resume_checkpoint_downloads_latest_and_links(tmp_path, monkeypatch):
+    from lerobot.common import train_utils
+
+    out = tmp_path / "run"
+
+    def fake_snapshot_download(repo_id, repo_type, allow_patterns, local_dir):
+        # Mimic the Hub layout the real download materializes locally.
+        assert allow_patterns == "checkpoints/020000/*"
+        (Path(local_dir) / "checkpoints" / "020000" / "pretrained_model").mkdir(parents=True)
+        return local_dir
+
+    monkeypatch.setattr("lerobot.common.train_utils.snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(
+        "lerobot.common.train_utils.find_latest_hub_checkpoint", lambda repo_id: "checkpoints/020000"
+    )
+
+    checkpoint_dir = train_utils.resolve_resume_checkpoint("u/run", out)
+
+    assert checkpoint_dir == out / CHECKPOINTS_DIR / "020000"
+    last = out / CHECKPOINTS_DIR / LAST_CHECKPOINT_LINK
+    assert last.is_symlink()
+    # `last` points at the downloaded step dir.
+    assert (last.parent / last.readlink()).resolve() == checkpoint_dir.resolve()
+
+
+def test_resolve_resume_checkpoint_raises_without_checkpoints(tmp_path, monkeypatch):
+    from lerobot.common import train_utils
+
+    monkeypatch.setattr("lerobot.common.train_utils.find_latest_hub_checkpoint", lambda repo_id: None)
+    with pytest.raises(FileNotFoundError, match="No checkpoint"):
+        train_utils.resolve_resume_checkpoint("u/run", tmp_path / "run")

tests/utils/test_visualization_utils.py

@@ -30,19 +30,25 @@ from lerobot.utils.constants import OBS_STATE
 @pytest.fixture
 def mock_rerun(monkeypatch):
     """
-    Provide a mock `rerun` module so tests don't depend on the real library.
-    Also reload the module-under-test so it binds to this mock `rr`.
+    Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
+    depend on the real library. Also reload the module-under-test so it binds to
+    this mock `rr`.
     """
     calls = []
+    blueprints = []
 
     class DummyScalar:
         def __init__(self, value):
-            self.value = float(value)
+            # Scalars may be built from a single float or from a 1D array batch.
+            self.value = value
 
     class DummyImage:
         def __init__(self, arr):
             self.arr = arr
 
+        def compress(self, *a, **k):
+            return self
+
     class DummyDepthImage:
         def __init__(self, arr, colormap=None):
             self.arr = arr
@@ -54,6 +60,21 @@ def mock_rerun(monkeypatch):
             obj = kwargs.pop("entity")
         calls.append((key, obj, kwargs))
 
+    def dummy_send_blueprint(blueprint, *a, **k):
+        blueprints.append(blueprint)
+
+    # Mock the `rerun.blueprint` submodule used to build the layout.
+    dummy_rrb = SimpleNamespace(
+        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
+            kind="Spatial2DView", origin=origin, name=name
+        ),
+        TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
+            kind="TimeSeriesView", name=name, contents=contents
+        ),
+        Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
+        Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
+    )
+
     dummy_rr = SimpleNamespace(
         __name__="rerun",
         __package__="rerun",
@@ -63,20 +84,23 @@ def mock_rerun(monkeypatch):
         DepthImage=DummyDepthImage,
         components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
         log=dummy_log,
+        send_blueprint=dummy_send_blueprint,
         init=lambda *a, **k: None,
         spawn=lambda *a, **k: None,
+        blueprint=dummy_rrb,
     )
 
-    # Inject fake module into sys.modules
+    # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
     monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
+    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
 
     # Now import and reload the module under test, to bind to our rerun mock
     import lerobot.utils.visualization_utils as vu
 
     importlib.reload(vu)
 
-    # Expose both the reloaded module and the call recorder
-    yield vu, calls
+    # Expose the reloaded module, the call recorder and the captured blueprints
+    yield vu, calls, blueprints
 
 
 def _keys(calls):
@@ -99,8 +123,13 @@ def _kwargs_for(calls, key):
     raise KeyError(f"Key {key} not found in calls: {calls}")
 
 
+def _views_by_kind(blueprint, kind):
+    """Return the views of a given kind from the (single) blueprint's grid."""
+    return [v for v in blueprint.root.views if v.kind == kind]
+
+
 def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
-    vu, calls = mock_rerun
+    vu, calls, blueprints = mock_rerun
 
     # Build EnvTransition dict
     obs = {
@@ -110,7 +139,7 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
     }
     act = {
         "action.throttle": 0.7,
-        # 1D array should log individual Scalars with suffix _i
+        # 1D array should be logged as a single Scalars batch under one entity path
         "action.vector": np.array([1.0, 2.0], dtype=np.float32),
     }
     transition = {
@@ -127,31 +156,28 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
     # - observation.state.temperature -> Scalars
     # - observation.camera -> Image (HWC) with static=True
     # - action.throttle -> Scalars
-    # - action.vector_0, action.vector_1 -> Scalars
+    # - action.vector -> single Scalars batch (no per-element suffix)
     expected_keys = {
         f"{OBS_STATE}.temperature",
         "observation.camera",
         "action.throttle",
-        "action.vector_0",
-        "action.vector_1",
+        "action.vector",
     }
     assert set(_keys(calls)) == expected_keys
 
     # Check scalar types and values
     temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
     assert type(temp_obj).__name__ == "DummyScalar"
-    assert temp_obj.value == pytest.approx(25.0)
+    assert float(temp_obj.value) == pytest.approx(25.0)
 
     throttle_obj = _obj_for(calls, "action.throttle")
     assert type(throttle_obj).__name__ == "DummyScalar"
-    assert throttle_obj.value == pytest.approx(0.7)
+    assert float(throttle_obj.value) == pytest.approx(0.7)
 
-    v0 = _obj_for(calls, "action.vector_0")
-    v1 = _obj_for(calls, "action.vector_1")
-    assert type(v0).__name__ == "DummyScalar"
-    assert type(v1).__name__ == "DummyScalar"
-    assert v0.value == pytest.approx(1.0)
-    assert v1.value == pytest.approx(2.0)
+    # 1D vector logged as a single batched Scalars under one entity path
+    vec = _obj_for(calls, "action.vector")
+    assert type(vec).__name__ == "DummyScalar"
+    np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
 
     # Check image handling: CHW -> HWC
     img_obj = _obj_for(calls, "observation.camera")
@@ -159,9 +185,24 @@ def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
     assert img_obj.arr.shape == (10, 20, 3)  # transposed
     assert _kwargs_for(calls, "observation.camera").get("static", False) is True  # static=True for images
 
+    # A blueprint should have been built and sent exactly once, and cached on the function.
+    assert len(blueprints) == 1
+    assert vu.log_rerun_data.blueprint is blueprints[0]
+
+    bp = blueprints[0]
+    # One spatial view per image path
+    spatial_views = _views_by_kind(bp, "Spatial2DView")
+    assert {v.origin for v in spatial_views} == {"observation.camera"}
+
+    # One time-series view each for observation and action scalars
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert set(ts_views) == {"observation", "action"}
+    assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
+    assert ts_views["action"].contents == ["action.throttle", "action.vector"]
+
 
 def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
-    vu, calls = mock_rerun
+    vu, calls, blueprints = mock_rerun
 
     # First dict without prefixes treated as observation
     # Second dict without prefixes treated as action
@@ -180,14 +221,12 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
     # First dict was treated as observation, second as action
     vu.log_rerun_data(observation=obs_plain, action=act_plain)
 
-    # Expected keys with auto-prefixes
+    # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
     expected = {
         "observation.temp",
         "observation.img",
         "action.throttle",
-        "action.vec_0",
-        "action.vec_1",
-        "action.vec_2",
+        "action.vec",
     }
     logged = set(_keys(calls))
     assert logged == expected
@@ -195,11 +234,11 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
     # Scalars
     t = _obj_for(calls, "observation.temp")
     assert type(t).__name__ == "DummyScalar"
-    assert t.value == pytest.approx(1.5)
+    assert float(t.value) == pytest.approx(1.5)
 
     throttle = _obj_for(calls, "action.throttle")
     assert type(throttle).__name__ == "DummyScalar"
-    assert throttle.value == pytest.approx(0.3)
+    assert float(throttle.value) == pytest.approx(0.3)
 
     # Image stays HWC
     img = _obj_for(calls, "observation.img")
@@ -207,15 +246,23 @@ def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
     assert img.arr.shape == (5, 6, 3)
     assert _kwargs_for(calls, "observation.img").get("static", False) is True
 
-    # Vectors
-    for i, val in enumerate([9, 8, 7]):
-        o = _obj_for(calls, f"action.vec_{i}")
-        assert type(o).__name__ == "DummyScalar"
-        assert o.value == pytest.approx(val)
+    # Vector logged as a single batched Scalars under one entity path
+    vec = _obj_for(calls, "action.vec")
+    assert type(vec).__name__ == "DummyScalar"
+    np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
+
+    # Blueprint sent once with the expected view layout
+    assert len(blueprints) == 1
+    bp = blueprints[0]
+    spatial_views = _views_by_kind(bp, "Spatial2DView")
+    assert {v.origin for v in spatial_views} == {"observation.img"}
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert ts_views["observation"].contents == ["observation.temp"]
+    assert ts_views["action"].contents == ["action.throttle", "action.vec"]
 
 
 def test_log_rerun_data_kwargs_only(mock_rerun):
-    vu, calls = mock_rerun
+    vu, calls, blueprints = mock_rerun
 
     vu.log_rerun_data(
         observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
@@ -229,7 +276,7 @@ def test_log_rerun_data_kwargs_only(mock_rerun):
 
     temp = _obj_for(calls, "observation.temp")
     assert type(temp).__name__ == "DummyScalar"
-    assert temp.value == pytest.approx(10.0)
+    assert float(temp.value) == pytest.approx(10.0)
 
     img = _obj_for(calls, "observation.gray")
     assert type(img).__name__ == "DummyDepthImage"  # single-channel -> DepthImage
@@ -238,4 +285,26 @@ def test_log_rerun_data_kwargs_only(mock_rerun):
 
     a = _obj_for(calls, "action.a")
     assert type(a).__name__ == "DummyScalar"
-    assert a.value == pytest.approx(1.0)
+    assert float(a.value) == pytest.approx(1.0)
+
+    # Blueprint sent once, with a spatial view for the image and time-series views for scalars
+    assert len(blueprints) == 1
+    bp = blueprints[0]
+    assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert ts_views["observation"].contents == ["observation.temp"]
+    assert ts_views["action"].contents == ["action.a"]
+
+
+def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
+    """The blueprint is built from the first call and not resent on subsequent calls."""
+    vu, calls, blueprints = mock_rerun
+
+    vu.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
+    assert len(blueprints) == 1
+    first_blueprint = vu.log_rerun_data.blueprint
+
+    vu.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
+    # Still only one blueprint, and the cached one is unchanged.
+    assert len(blueprints) == 1
+    assert vu.log_rerun_data.blueprint is first_blueprint

uv.lock

@@ -3395,7 +3395,7 @@ requires-dist = [
     { name = "qwen-vl-utils", marker = "extra == 'qwen-vl-utils-dep'", specifier = ">=0.0.11,<0.1.0" },
     { name = "reachy2-sdk", marker = "extra == 'reachy2'", specifier = ">=1.0.15,<1.1.0" },
     { name = "requests", specifier = ">=2.32.0,<3.0.0" },
-    { name = "rerun-sdk", marker = "extra == 'viz'", specifier = ">=0.24.0,<0.27.0" },
+    { name = "rerun-sdk", marker = "extra == 'viz'", specifier = ">=0.24.0,<0.34.0" },
     { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.14.1" },
     { name = "safetensors", specifier = ">=0.4.3,<1.0.0" },
     { name = "scikit-image", marker = "extra == 'video-benchmark'", specifier = ">=0.23.2,<0.26.0" },
@@ -5803,21 +5803,21 @@ wheels = [
 
 [[package]]
 name = "rerun-sdk"
-version = "0.26.2"
+version = "0.33.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "attrs" },
     { name = "numpy" },
     { name = "pillow" },
+    { name = "psutil" },
     { name = "pyarrow" },
     { name = "typing-extensions" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/4b/4a/767c20e1529d74d9be5b5e55c6c26b63a6918ef3c1709fc422d08a460114/rerun_sdk-0.26.2-cp39-abi3-macosx_10_12_x86_64.whl", hash = "sha256:3d4151c9a3484e112b53d1df90c8fa07397dc7b8bfbb420f09e011eff20f1ef2", size = 93349439, upload-time = "2025-10-27T11:34:10.745Z" },
-    { url = "https://files.pythonhosted.org/packages/2b/3d/d8dd0af9c287a85d51ec99d69406cc4b94a9feb1d6f192d3bbcaac9f0b81/rerun_sdk-0.26.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:03977d2aba4966d9a70b682eca196123fda11408fecd733441ede9916c6341e2", size = 86323042, upload-time = "2025-10-27T11:34:17.995Z" },
-    { url = "https://files.pythonhosted.org/packages/13/29/53d8d98799ab32418fd4ba6834d6a5749c31f56160d3c87f52a7219887e9/rerun_sdk-0.26.2-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:b6128c3c4f014cae5be18e4d37657c5932d1bcdb2ce5e9d4b488a6eed47f7437", size = 92677274, upload-time = "2025-10-27T11:34:22.601Z" },
-    { url = "https://files.pythonhosted.org/packages/f5/86/0b9c8f56398b4fc85f8e99279907c258413a297e5603f8f2537fe5806e51/rerun_sdk-0.26.2-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a6f97b60aaa7d4e8c6124a3f6b97ce9dbd09520050955f0e0bdacb72b0eb106a", size = 98768129, upload-time = "2025-10-27T11:34:27.36Z" },
-    { url = "https://files.pythonhosted.org/packages/be/e7/99fc91c0f99f69d7d43e1db0a6f6cb8273ffc02111539bfc1fee43749bad/rerun_sdk-0.26.2-cp39-abi3-win_amd64.whl", hash = "sha256:a493ad6c8357022cba2ca6f8954a81d0faf984b0b22154eb1d976bfc7649df63", size = 84267089, upload-time = "2025-10-27T11:34:32.023Z" },
+    { url = "https://files.pythonhosted.org/packages/16/07/380198590b194f1c17052d672865aa4a56e606eae47665f66edfb391999d/rerun_sdk-0.33.1-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c0115b710289022587bd2e9ecb715f98d1e87dd7d8ac48e053324131d4addc89", size = 125706253, upload-time = "2026-06-22T09:04:18.974Z" },
+    { url = "https://files.pythonhosted.org/packages/f8/eb/6741bbf6868175ab126aff58d372066241c6cd2fc1c4f82ed64069728e73/rerun_sdk-0.33.1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:14451d31bc7bd0f7c6bcc9c1213ed679ab81b65ecd8b36eae99f738219897dc5", size = 135278374, upload-time = "2026-06-22T09:04:26.768Z" },
+    { url = "https://files.pythonhosted.org/packages/04/28/fd3b832652900fe3415739e7411c8af8af4c44e9e1a9d55d79e37f7f9094/rerun_sdk-0.33.1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:0f87da7a270614074aca37846350f9e257f65081345474748f578c7da64fdeba", size = 139565470, upload-time = "2026-06-22T09:04:34.941Z" },
+    { url = "https://files.pythonhosted.org/packages/19/41/339920f5a6734054c07bcae543365a7ef3368ceee3eb67906e2e38bd1d67/rerun_sdk-0.33.1-cp310-abi3-win_amd64.whl", hash = "sha256:b2c2af67f3c2a85b282669d97b52596593fcfdd19bf57c423f18827c837a6e49", size = 120411717, upload-time = "2026-06-22T09:04:42.643Z" },
 ]
 
 [[package]]