Admit local Diffusion, VQ-BeT and TDMPC recipes; refuse unsupported devices

docs/local-policy-recipes.md

@@ -0,0 +1,41 @@
+# Local policy recipes
+
+The local LeRobot v3 input contract admits `diffusion`, `vqbet` and `tdmpc`
+with `mode=full`, `phase=train`, an immutable dataset content digest and explicit
+recipe settings. ACT and pretrained SmolVLA retain their existing contracts.
+Admission does not establish task quality, export parity or robot support.
+Use `skip_export` for families without an exporter.
+
+Diffusion cameras must share their RGB shape. Each side must be at least 32
+pixels, with at least one side larger than 32 for ResNet18 initialization.
+VQ-BeT and TDMPC accept exactly one RGB camera of at least 64 pixels per side;
+TDMPC requires square images. State and action vectors come from the dataset.
+Unsupported observation fields are refused explicitly.
+Local VQ-BeT requires explicit `policy.device=cpu`: LeRobot 0.6.1's factory
+refuses MPS. Preflight returns that reason before starting the trainer.
+
+TDMPC additionally needs `next.reward`, logical shape `[1]`, stored as a finite
+scalar float in every transition. Singleton list storage is refused because
+native TDMPC expects a scalar reward batch. This recipe performs offline
+world/reward/Q learning; it does not implement SAC's online actor/learner loop.
+VQ-BeT needs a completed VQ codebook phase before inference.
+
+Local resume for these three families verifies the dataset again and selects a
+complete numeric checkpoint. Saved LeRobot 0.6.1 config, trainer step, optimizer,
+RNG and scheduler where applicable are required. Family, dataset root, step
+budget, batch, seed and explicit LeRobot overrides must match. Changed settings,
+old 0.5.1 configs, changed output directories or incomplete state return named refusals. The caller must
+retain the checkpoint directory and dataset. Exact numerical equality with an
+uninterrupted run is unverified.
+
+Studio's adapter adds atomic 500-step saves, latest-two retention, checkpoint
+hash review and CPU/MPS settings. Those are Studio responsibilities rather than
+new Tether defaults. Local provider charges are $0. No hardware or paid-service
+qualification follows from these recipes.
+
+Verified in simulation on 10/15/2026: native local dataset VQ-BeT and TDMPC CPU
+four-step plumbing and saved checkpoints; TDMPC uses explicitly synthetic zero
+rewards for that probe. Full Tether suite: 4,224 passed, 140 skipped with the
+S-11 interpreter and supplied MCP/async dependencies. Four native TDMPC MPS
+updates also produced finite losses and weights. VQ-BeT MPS is unavailable;
+hardware and learned task quality remain unverified.

src/tether/finetune/local_dataset.py

@@ -746,19 +746,20 @@ def validate_local_config(cfg: Any) -> None:
         "dataset_root is required with a valid optional dataset_manifest_sha256 or dataset_content_sha256.",
     )
     _require(
-        not cfg.skip_preflight and not cfg.resume,
+        not cfg.skip_preflight and (not cfg.resume or
+            (raw and cfg.policy in {"diffusion", "vqbet", "tdmpc"} and cfg.mode == "full")),
         "local-config",
-        "Local dataset admission does not support skip_preflight or resume.",
+        "Local dataset admission does not support skip_preflight or resume outside the state-bound local Diffusion/VQ-BeT/TDMPC recipes.",
     )
     _require(
         cfg.phase == "train"
         and cfg.backend == "lerobot"
-        and ((cfg.policy == "act" and cfg.mode == "full") or (cfg.policy == "auto" and cfg.mode == "lora" and bool(cfg.base)))
+        and ((cfg.policy in {"act", "diffusion", "vqbet", "tdmpc"} and cfg.mode == "full") or (cfg.policy == "auto" and cfg.mode == "lora" and bool(cfg.base)))
         and (raw or cfg.policy == "auto")
         and bool(cfg.dataset)
         and not cfg.base_dataset,
         "unsupported-profile",
-        "Local datasets require ACT from scratch (policy=act, mode=full) or pretrained SmolVLA LoRA (policy=auto, mode=lora), phase=train.",
+        "Local datasets require ACT, Diffusion, VQ-BeT or TDMPC from scratch (policy=act|diffusion|vqbet|tdmpc, mode=full) or pretrained SmolVLA LoRA (policy=auto, mode=lora), phase=train.",
     )
     _require(not raw or not (Path(root).expanduser().resolve().parent / "export-receipt.json").exists(), "local-config", "dataset_manifest_sha256: a Studio export receipt requires its explicit digest; it cannot be downgraded to raw admission.")
     export = Path(root).expanduser().resolve()

src/tether/finetune/preflight/schema.py

@@ -143,8 +143,49 @@ def check_schema(cfg: FinetuneConfig) -> PreflightCheck:
         # its pretraining embodiment; the command clears them
         # (--policy.input_features=null) so LeRobot derives them from this
         # dataset. Only the padding bounds bind, as on the export path.
-        if cfg.policy == "act" and cfg.mode == "full":
-            return PreflightCheck("schema", "ok", "Local features match the ACT recipe.", {"features": features})
+        if cfg.policy in {"act", "diffusion", "vqbet", "tdmpc"} and cfg.mode == "full":
+            if cfg.policy != "act":
+                cameras = {key: feature["shape"] for key, feature in features.items()
+                           if feature["dtype"] in {"image", "video"}}
+                if not cameras:
+                    return PreflightCheck("schema", "fail", "features.observation.images: local image policy requires a camera")
+                unsupported = [key for key in features if key.startswith("observation.")
+                               and key != "observation.state" and key not in cameras]
+                if unsupported:
+                    return PreflightCheck("schema", "fail", f"features.{unsupported[0]}: this recipe consumes state and camera observations only")
+            if cfg.policy == "diffusion":
+                cameras = {key: feature["shape"] for key, feature in features.items()
+                           if feature["dtype"] in {"image", "video"}}
+                first = next(iter(cameras.values()))
+                for key, shape in cameras.items():
+                    if shape != first:
+                        return PreflightCheck("schema", "fail", f"features.{key}.shape: Diffusion cameras must share an RGB shape")
+                    if min(shape[:2]) < 32 or shape[:2] == [32, 32]:
+                        return PreflightCheck("schema", "fail", f"features.{key}.shape: Diffusion requires >=32 pixel sides and a ResNet18 initialization larger than 32 x 32")
+            if cfg.policy in {"vqbet", "tdmpc"}:
+                if cfg.policy == "vqbet" and cfg.extra_lerobot_args.get("policy.device") != "cpu":
+                    return PreflightCheck("schema", "fail", "policy.device: this local VQ-BeT recipe requires explicit CPU; LeRobot 0.6.1 refuses MPS")
+                cameras = {key: f["shape"] for key, f in features.items() if f["dtype"] in {"image", "video"}}
+                if len(cameras) != 1:
+                    return PreflightCheck("schema", "fail", f"features.observation.images: {cfg.policy} requires exactly one camera")
+                key, shape = next(iter(cameras.items()))
+                if min(shape[:2]) < 64 or (cfg.policy == "tdmpc" and shape[0] != shape[1]):
+                    return PreflightCheck("schema", "fail", f"features.{key}.shape: requires >=64 pixel sides; TDMPC requires square images")
+                if cfg.policy == "tdmpc":
+                    reward = features.get("next.reward", {})
+                    if reward.get("dtype") not in {"float32", "float64"} or reward.get("shape") != [1]:
+                        return PreflightCheck("schema", "fail", "features.next.reward: TDMPC requires finite scalar transition rewards with shape [1]")
+                    import math
+                    import pyarrow.parquet as pq
+                    for path in sorted(Path(cfg.dataset_root).glob("data/**/*.parquet")):
+                        parquet = pq.ParquetFile(path)
+                        if "next.reward" not in parquet.schema_arrow.names:
+                            return PreflightCheck("schema", "fail", f"features.next.reward: missing stored column in {path.name}")
+                        for batch in parquet.iter_batches(batch_size=1024, columns=["next.reward"]):
+                            for value in batch.column(0).to_pylist():
+                                if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
+                                    return PreflightCheck("schema", "fail", f"features.next.reward: invalid scalar reward in {path.name}")
+            return PreflightCheck("schema", "ok", f"Local features match the {cfg.policy.upper()} recipe.", {"features": features})
         if not isinstance(base_config, dict) or base_config.get("type") != "smolvla":
             return PreflightCheck("schema", "fail", "base.config.type: requires a resolvable supported SmolVLA checkpoint")
         for key, maximum in (("action", "max_action_dim"), ("observation.state", "max_state_dim")):

src/tether/finetune/run.py

@@ -190,6 +190,48 @@ def _build_lerobot_command(cfg: FinetuneConfig) -> list[str]:
     # uploading). We don't push to Hub, but the validator runs anyway.
     # Pass a placeholder derived from the output dir name.
     # Customers who want to auto-push can override via extra_lerobot_args.
+    if cfg.resume and local_export is not None:
+        import json
+        candidates = sorted(p for p in (lerobot_output / "checkpoints").glob("*")
+                            if p.is_dir() and not p.is_symlink() and p.name.isdigit())
+        for checkpoint in reversed(candidates):
+            required = ("training_step.json", "optimizer_state.safetensors", "optimizer_param_groups.json", "rng_state.safetensors")
+            if cfg.policy != "tdmpc":
+                required += ("scheduler_state.json",)
+            if not all((checkpoint / "training_state" / name).is_file() for name in required):
+                continue
+            config_path = checkpoint / "pretrained_model/train_config.json"
+            if not config_path.is_file():
+                continue
+            saved = json.loads(config_path.read_text())
+            if "eval_freq" in saved or "env_eval_freq" not in saved:
+                raise ValueError("resume.lerobot_version: requires a LeRobot 0.6.1 trainer config, not a 0.5.1 checkpoint")
+            recorded_step = json.loads((checkpoint / "training_state/training_step.json").read_text()).get("step")
+            if type(recorded_step) is not int or recorded_step != int(checkpoint.name):
+                raise ValueError("resume.training_step: checkpoint folder differs from trainer state")
+            if saved.get("policy", {}).get("type") != cfg.policy:
+                raise ValueError("resume.policy.type: checkpoint family differs from reviewed recipe")
+            if Path(saved["dataset"]["root"]).resolve() != Path(cfg.dataset_root).resolve():
+                raise ValueError("resume.dataset.root: checkpoint dataset differs from immutable snapshot")
+            if not saved.get("output_dir") or Path(saved["output_dir"]).resolve() != lerobot_output.resolve():
+                raise ValueError("resume.output_dir: checkpoint differs from the reviewed job directory")
+            for key, expected in (("steps", cfg.num_steps), ("batch_size", cfg.batch_size), ("seed", cfg.seed)):
+                if saved.get(key) != expected:
+                    raise ValueError(f"resume.{key}: checkpoint differs from reviewed settings")
+            for key, expected in cfg.extra_lerobot_args.items():
+                if isinstance(expected, str):
+                    try:
+                        expected = json.loads(expected)
+                    except json.JSONDecodeError:
+                        pass
+                actual = saved
+                for part in key.split("."):
+                    actual = actual.get(part) if isinstance(actual, dict) else None
+                if actual != expected:
+                    raise ValueError(f"resume.{key}: checkpoint differs from reviewed settings")
+            return ["lerobot-train", f"--config_path={config_path}", "--resume=true"]
+        raise ValueError("resume.checkpoint: no complete numeric trainer checkpoint; retain previous checkpoint and retry")
+
     repo_id = f"local/{cfg.output.name}"
 
     cmd = [
@@ -217,7 +259,7 @@ def _build_lerobot_command(cfg: FinetuneConfig) -> list[str]:
         cmd.append(f"--dataset.root={local_export['dataset_root']}")
     if cfg.dataset_revision:
         cmd.append(f"--dataset.revision={cfg.dataset_revision}")
-    if is_from_scratch and cfg.chunk_size:
+    if is_from_scratch and explicit_policy == "act" and cfg.chunk_size:
         # ACT (and similar chunked policies) need chunk_size; pretrained bases
         # bake this in. auto_soarm convention (per its train.py): set
         # n_action_steps == chunk_size so ACTConfig.__post_init__ doesn't

tests/test_s38_diffusion_local.py

@@ -0,0 +1,116 @@
+"""Diffusion joins raw v3 admission; no relaxation for exports or other families."""
+from pathlib import Path
+import pytest
+from tether.finetune.config import FinetuneConfig
+from tether.finetune.local_dataset import validate_local_config, LocalDatasetError
+from tether.finetune.run import _build_lerobot_command
+
+
+def config(tmp_path, **changes):
+    return FinetuneConfig(base='', policy='diffusion', mode='full', dataset='studio-local/own',
+        dataset_root=tmp_path/'data', dataset_content_sha256='a'*64, output=tmp_path/'out', **changes)
+
+
+def test_raw_diffusion_admission_and_no_act_flags(tmp_path):
+    cfg = config(tmp_path)
+    validate_local_config(cfg)
+    # Command schema without local verification; full verified-byte path is covered by the S-38 real smoke.
+    cfg.dataset_root = None; cfg.dataset_content_sha256 = None
+    command = _build_lerobot_command(cfg)
+    assert '--policy.type=diffusion' in command
+    assert not any(arg.startswith('--policy.chunk_size=') for arg in command)
+
+@pytest.mark.parametrize('family', ['sac', 'smolvla', 'wall_x', 'multi_task_dit', 'gaussian_actor'])
+def test_unadmitted_family_stays_refused(tmp_path, family):
+    cfg = config(tmp_path); cfg.policy = family
+    with pytest.raises(LocalDatasetError, match='from scratch'):
+        validate_local_config(cfg)
+
+
+def test_diffusion_does_not_admit_the_smolvla_export_profile(tmp_path):
+    cfg = config(tmp_path); cfg.dataset_content_sha256 = None; cfg.dataset_manifest_sha256 = 'b'*64
+    with pytest.raises(LocalDatasetError):
+        validate_local_config(cfg)
+
+
+@pytest.mark.parametrize('family', ['vqbet', 'tdmpc'])
+def test_local_offline_family_admitted_without_act_chunk(tmp_path, family):
+    cfg = config(tmp_path); cfg.policy = family
+    validate_local_config(cfg)
+    cfg.dataset_root = None; cfg.dataset_content_sha256 = None
+    command = _build_lerobot_command(cfg)
+    assert f'--policy.type={family}' in command
+    assert not any(arg.startswith('--policy.chunk_size=') for arg in command)
+
+
+def test_local_resume_selects_complete_numeric_checkpoint_and_binds_input(tmp_path, monkeypatch):
+    import json
+    import tether.finetune.lerobot_v3 as v3
+    import tether.finetune.preflight.schema as schema
+    import tether.finetune.run as run
+    cfg = config(tmp_path); cfg.resume = True
+    saved = cfg.output/'training/checkpoints/000500'
+    (saved/'pretrained_model').mkdir(parents=True)
+    (saved/'training_state').mkdir()
+    for name in ('training_step.json','optimizer_state.safetensors','optimizer_param_groups.json','rng_state.safetensors','scheduler_state.json'):
+        (saved/'training_state'/name).write_text('{}')
+    (saved/'training_state/training_step.json').write_text('{"step":500}')
+    data = {'env_eval_freq':0,'policy': {'type':'diffusion'}, 'dataset': {'root':str(cfg.dataset_root)}, 'output_dir':str(cfg.output/'training'),
+            'steps':cfg.num_steps,'batch_size':cfg.batch_size,'seed':cfg.seed}
+    path = saved/'pretrained_model/train_config.json';path.write_text(json.dumps(data))
+    monkeypatch.setattr(run,'verify_local_dataset', lambda cfg: {'dataset_root':str(cfg.dataset_root)})
+    from tether.finetune.preflight.result import PreflightCheck
+    monkeypatch.setattr(schema,'check_schema',lambda cfg:PreflightCheck('schema','ok','fixture'))
+    command = _build_lerobot_command(cfg)
+    assert command == ['lerobot-train', f'--config_path={path}', '--resume=true']
+    data['batch_size'] += 1;path.write_text(json.dumps(data))
+    with pytest.raises(ValueError,match='resume.batch_size'):_build_lerobot_command(cfg)
+    data['batch_size'] -= 1; data['policy']['horizon']=16;path.write_text(json.dumps(data))
+    cfg.extra_lerobot_args={'policy.horizon':32}
+    with pytest.raises(ValueError,match='resume.policy.horizon'):_build_lerobot_command(cfg)
+    cfg.extra_lerobot_args={'policy.horizon':16}
+    assert _build_lerobot_command(cfg)==['lerobot-train',f'--config_path={path}','--resume=true']
+    data['output_dir']=str(tmp_path/'elsewhere');path.write_text(json.dumps(data))
+    with pytest.raises(ValueError,match='resume.output_dir'):_build_lerobot_command(cfg)
+
+
+@pytest.mark.parametrize('reward,ok', [(0.,True),([1.],False),(float('nan'),False)])
+def test_tdmpc_refuses_non_scalar_or_nonfinite_stored_rewards(tmp_path,monkeypatch,reward,ok):
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import tether.finetune.preflight.schema as schema
+    cfg=config(tmp_path);cfg.policy='tdmpc'
+    data=cfg.dataset_root/'data/chunk-000/file-000.parquet';data.parent.mkdir(parents=True)
+    pq.write_table(pa.table({'next.reward':[reward]}),data)
+    monkeypatch.setattr(schema,'verify_local_dataset',lambda cfg:{'info':{'features':{
+        'observation.state':{'dtype':'float32','shape':[6]},
+        'action':{'dtype':'float32','shape':[6]},
+        'observation.images.wrist':{'dtype':'image','shape':[96,96,3]},
+        'next.reward':{'dtype':'float32','shape':[1]}}}})
+    result=schema.check_schema(cfg)
+    assert result.severity==('ok' if ok else 'fail')
+    if not ok:assert 'features.next.reward' in result.summary
+
+
+def test_local_vqbet_requires_explicit_cpu_before_native_factory(tmp_path,monkeypatch):
+    import tether.finetune.preflight.schema as schema
+    cfg=config(tmp_path);cfg.policy='vqbet';cfg.extra_lerobot_args={'policy.device':'mps'}
+    monkeypatch.setattr(schema,'verify_local_dataset',lambda cfg:{'info':{'features':{
+        'observation.state':{'dtype':'float32','shape':[6]},'action':{'dtype':'float32','shape':[6]},
+        'observation.images.wrist':{'dtype':'image','shape':[96,96,3]}}}})
+    result=schema.check_schema(cfg)
+    assert result.is_blocking and 'policy.device' in result.summary and 'refuses MPS' in result.summary
+
+
+def test_tdmpc_missing_stored_reward_is_named(tmp_path,monkeypatch):
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import tether.finetune.preflight.schema as schema
+    cfg=config(tmp_path);cfg.policy='tdmpc'
+    path=cfg.dataset_root/'data/chunk-000/file-000.parquet';path.parent.mkdir(parents=True)
+    pq.write_table(pa.table({'index':[0]}),path)
+    monkeypatch.setattr(schema,'verify_local_dataset',lambda cfg:{'info':{'features':{
+        'observation.state':{'dtype':'float32','shape':[6]},'action':{'dtype':'float32','shape':[6]},
+        'observation.images.wrist':{'dtype':'image','shape':[96,96,3]},'next.reward':{'dtype':'float32','shape':[1]}}}})
+    result=schema.check_schema(cfg)
+    assert result.is_blocking and 'features.next.reward: missing stored column' in result.summary