diff --git a/docs/local-policy-recipes.md b/docs/local-policy-recipes.md new file mode 100644 index 00000000..5e8bcba7 --- /dev/null +++ b/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. diff --git a/src/tether/finetune/local_dataset.py b/src/tether/finetune/local_dataset.py index be50f0db..ee823160 100644 --- a/src/tether/finetune/local_dataset.py +++ b/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() diff --git a/src/tether/finetune/preflight/schema.py b/src/tether/finetune/preflight/schema.py index 8ecb5f50..8663ac5f 100644 --- a/src/tether/finetune/preflight/schema.py +++ b/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")): diff --git a/src/tether/finetune/run.py b/src/tether/finetune/run.py index 81bfca93..cfcf67aa 100644 --- a/src/tether/finetune/run.py +++ b/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 diff --git a/tests/test_s38_diffusion_local.py b/tests/test_s38_diffusion_local.py new file mode 100644 index 00000000..ab3e5a43 --- /dev/null +++ b/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