diff --git a/frontend/src/components/training/ConfigurationTab.tsx b/frontend/src/components/training/ConfigurationTab.tsx index 8106bb79..caa0ab8c 100644 --- a/frontend/src/components/training/ConfigurationTab.tsx +++ b/frontend/src/components/training/ConfigurationTab.tsx @@ -1,5 +1,6 @@ import React from 'react'; import EssentialsCard from './config/EssentialsCard'; +import GrootCard from './config/GrootCard'; import AdvancedCard from './config/AdvancedCard'; import TargetCard from './config/TargetCard'; import { ConfigComponentProps } from './types'; @@ -38,6 +39,7 @@ const ConfigurationTab: React.FC = ({ datasets={datasets} datasetsLoading={datasetsLoading} /> + ); diff --git a/frontend/src/components/training/config/EssentialsCard.tsx b/frontend/src/components/training/config/EssentialsCard.tsx index b99397ee..fac6b77d 100644 --- a/frontend/src/components/training/config/EssentialsCard.tsx +++ b/frontend/src/components/training/config/EssentialsCard.tsx @@ -92,6 +92,7 @@ const EssentialsCard: React.FC = ({ config, updateConfig, d Diffusion Policy PI0 SmolVLA + GR00T N1.7 TD-MPC VQ-BeT PI0 Fast diff --git a/frontend/src/components/training/config/GrootCard.tsx b/frontend/src/components/training/config/GrootCard.tsx new file mode 100644 index 00000000..ea820f0e --- /dev/null +++ b/frontend/src/components/training/config/GrootCard.tsx @@ -0,0 +1,178 @@ +import React, { useEffect, useRef } from 'react'; +import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card'; +import { Input } from '@/components/ui/input'; +import { NumberInput } from '@/components/ui/number-input'; +import { Label } from '@/components/ui/label'; +import { Switch } from '@/components/ui/switch'; +import { ConfigComponentProps } from '../types'; + +// Sensible GR00T N1.7 defaults, seeded the first time the user picks the policy +// so the form matches the canonical fine-tuning command out of the box. +const GROOT_DEFAULTS = { + policy_base_model_path: 'nvidia/GR00T-N1.7-3B', + policy_embodiment_tag: 'new_embodiment', + policy_chunk_size: 16, + policy_n_action_steps: 16, + policy_use_relative_actions: true, + policy_relative_exclude_joints: ['gripper'], + policy_use_bf16: true, + dataset_image_transforms_enable: true, +} as const; + +const GrootCard: React.FC = ({ config, updateConfig }) => { + const isGroot = config.policy_type === 'groot'; + const seededDefaults = useRef(false); + + // Seed defaults once when groot is selected. Each key is only written when + // still undefined, so the effect can't loop and never clobbers user edits. + useEffect(() => { + if (!isGroot || seededDefaults.current) return; + seededDefaults.current = true; + (Object.keys(GROOT_DEFAULTS) as (keyof typeof GROOT_DEFAULTS)[]).forEach((key) => { + if (config[key] === undefined) { + updateConfig(key, GROOT_DEFAULTS[key] as never); + } + }); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, [isGroot]); + + if (!isGroot) return null; + + const excludeJoints = (config.policy_relative_exclude_joints ?? []).join(', '); + + return ( + + + GR00T N1.7 + + +
+ + + updateConfig('policy_base_model_path', e.target.value || undefined) + } + placeholder="nvidia/GR00T-N1.7-3B" + className="bg-slate-900 border-slate-600 text-white rounded-lg" + /> +

+ HuggingFace repo of the pretrained GR00T backbone to fine-tune. +

+
+ +
+
+ + + updateConfig('policy_embodiment_tag', e.target.value || undefined) + } + placeholder="new_embodiment" + className="bg-slate-900 border-slate-600 text-white rounded-lg" + /> +
+ +
+ + + updateConfig( + 'policy_relative_exclude_joints', + e.target.value + .split(',') + .map((s) => s.trim()) + .filter(Boolean), + ) + } + placeholder="gripper" + className="bg-slate-900 border-slate-600 text-white rounded-lg" + /> +

+ Comma-separated joints kept absolute when relative actions are on. +

+
+ +
+ + updateConfig('policy_chunk_size', v)} + className="bg-slate-900 border-slate-600 text-white rounded-lg" + /> +
+ +
+ + updateConfig('policy_n_action_steps', v)} + className="bg-slate-900 border-slate-600 text-white rounded-lg" + /> +
+
+ +
+
+ + updateConfig('policy_use_relative_actions', checked) + } + className="data-[state=checked]:bg-green-500" + /> + +
+ +
+ updateConfig('policy_use_bf16', checked)} + className="data-[state=checked]:bg-green-500" + /> + +
+ +
+ + updateConfig('dataset_image_transforms_enable', checked) + } + className="data-[state=checked]:bg-green-500" + /> + +
+
+ + + ); +}; + +export default GrootCard; diff --git a/frontend/src/components/training/types.ts b/frontend/src/components/training/types.ts index d0ee5f44..cd4882dd 100644 --- a/frontend/src/components/training/types.ts +++ b/frontend/src/components/training/types.ts @@ -41,6 +41,16 @@ export interface TrainingConfig { // Advanced configuration use_policy_training_preset: boolean; + + // GR00T-specific configuration (only used when policy_type === "groot"). + dataset_image_transforms_enable?: boolean; + policy_base_model_path?: string; + policy_embodiment_tag?: string; + policy_chunk_size?: number; + policy_n_action_steps?: number; + policy_use_relative_actions?: boolean; + policy_relative_exclude_joints?: string[]; + policy_use_bf16?: boolean; } export interface TrainingStatus { diff --git a/frontend/src/lib/jobsApi.ts b/frontend/src/lib/jobsApi.ts index 088b727b..8cd82f70 100644 --- a/frontend/src/lib/jobsApi.ts +++ b/frontend/src/lib/jobsApi.ts @@ -49,6 +49,16 @@ export interface TrainingRequest { optimizer_weight_decay?: number; optimizer_grad_clip_norm?: number; use_policy_training_preset: boolean; + // GR00T-specific (only sent when policy_type === "groot"); backend ignores + // the policy_* ones for other policies. + dataset_image_transforms_enable?: boolean; + policy_base_model_path?: string; + policy_embodiment_tag?: string; + policy_chunk_size?: number; + policy_n_action_steps?: number; + policy_use_relative_actions?: boolean; + policy_relative_exclude_joints?: string[]; + policy_use_bf16?: boolean; // Optional target for runner dispatch; omitted ⇒ local. target?: { runner: "local" | "hf_cloud"; flavor?: string }; } diff --git a/frontend/src/pages/Training.tsx b/frontend/src/pages/Training.tsx index a361d042..ca15ae0f 100644 --- a/frontend/src/pages/Training.tsx +++ b/frontend/src/pages/Training.tsx @@ -93,6 +93,20 @@ function configToRequest(c: TrainingConfig): TrainingRequest { optimizer_weight_decay: c.optimizer_weight_decay, optimizer_grad_clip_norm: c.optimizer_grad_clip_norm, use_policy_training_preset: c.use_policy_training_preset, + // GR00T-specific. Only forward the policy_* fields for groot so other + // policies never receive flags their config doesn't define. + dataset_image_transforms_enable: c.dataset_image_transforms_enable, + ...(c.policy_type === "groot" + ? { + policy_base_model_path: c.policy_base_model_path, + policy_embodiment_tag: c.policy_embodiment_tag, + policy_chunk_size: c.policy_chunk_size, + policy_n_action_steps: c.policy_n_action_steps, + policy_use_relative_actions: c.policy_use_relative_actions, + policy_relative_exclude_joints: c.policy_relative_exclude_joints, + policy_use_bf16: c.policy_use_bf16, + } + : {}), }; } diff --git a/lelab/rollout.py b/lelab/rollout.py index cfff79a8..fa02e507 100644 --- a/lelab/rollout.py +++ b/lelab/rollout.py @@ -24,6 +24,7 @@ from __future__ import annotations import contextlib +import json import logging import os import re @@ -139,6 +140,53 @@ def _resolve_policy_path(policy_ref: str) -> str: raise ValueError(f"Unrecognised policy ref: {policy_ref!r}") +def _read_policy_config(policy_path: str) -> dict[str, Any]: + """Load pretrained_model/config.json if present.""" + config_path = Path(policy_path) / "config.json" + if not config_path.is_file(): + return {} + try: + with open(config_path) as fh: + data = json.load(fh) + return data if isinstance(data, dict) else {} + except Exception as exc: + logger.warning("Failed to read policy config at %s: %s", config_path, exc) + return {} + + +def _rollout_inference_args(policy_path: str) -> list[str]: + """Return extra lerobot-rollout flags for policies that reject sync inference. + + GR00T and other relative-action policies decode action chunks against the + observation state at inference time. The default sync backend calls + ``select_action`` per tick and can re-decode cached relative actions + against newer states, so lerobot requires the RTC/chunked rollout path + instead. + + GR00T also needs tuned RTC queue settings: the lerobot default + ``queue_threshold=30`` makes the background thread replan while a 16-step + chunk is still playing, which replaces the queue mid-motion and feels + like one action then a full recalculation.""" + cfg = _read_policy_config(policy_path) + policy_type = cfg.get("type") + needs_rtc = bool(cfg.get("use_relative_actions")) or policy_type == "groot" + if not needs_rtc: + return [] + + args = ["--inference.type=rtc"] + + n_action_steps = cfg.get("n_action_steps") + if isinstance(n_action_steps, int) and n_action_steps > 0: + args.append(f"--inference.rtc.execution_horizon={n_action_steps}") + elif policy_type == "groot": + args.append("--inference.rtc.execution_horizon=16") + + # Wait until the current chunk is consumed before replanning. GROOT docs + # recommend keeping this at 0 (never > 5) for stable real-robot rollout. + args.append("--inference.queue_threshold=0") + return args + + def _format_cameras_arg(cameras: dict[str, dict[str, Any]]) -> str: """Convert {name: {type, camera_index, width, height, fps}} into lerobot's CLI dict syntax. The frontend key `camera_index` is @@ -211,6 +259,11 @@ def _friendly_hint(error_text: str | None) -> str | None: return "A camera doesn't support the configured resolution — open camera settings and click Auto." if "permission" in low and ("port" in low or "com" in low): return "Couldn't open the serial port — close anything else using it, or run `lelab --stop`." + if "relative-action" in low or "relative chunk actions" in low: + return ( + "This policy was trained with relative actions and needs RTC (chunked) inference. " + "Update lelab and try again — recent versions enable that automatically." + ) return None @@ -286,6 +339,7 @@ def handle_start_inference(request: InferenceRequest) -> dict[str, Any]: f"--robot.id={follower_id}", f"--task={request.task}", f"--duration={request.duration_s}", + *_rollout_inference_args(policy_path), ] if request.cameras: cmd.append(f"--robot.cameras={_format_cameras_arg(request.cameras)}") diff --git a/lelab/train.py b/lelab/train.py index bfdc0226..257edeaf 100644 --- a/lelab/train.py +++ b/lelab/train.py @@ -18,6 +18,7 @@ lives in app/jobs.py. """ +import json import re from typing import TYPE_CHECKING @@ -35,6 +36,7 @@ class TrainingRequest(BaseModel): dataset_revision: str | None = None dataset_root: str | None = None dataset_episodes: list[int] | None = None + dataset_image_transforms_enable: bool = False # Policy configuration policy_type: str = "act" @@ -79,6 +81,17 @@ class TrainingRequest(BaseModel): policy_push_to_hub: bool = False policy_repo_id: str | None = None + # GR00T-specific policy options (only emitted when policy_type == "groot"; + # these flags don't exist on other policy configs, so draccus would reject + # them). All optional — unset fields fall back to the groot config defaults. + policy_base_model_path: str | None = None + policy_embodiment_tag: str | None = None + policy_chunk_size: int | None = None + policy_n_action_steps: int | None = None + policy_use_relative_actions: bool | None = None + policy_relative_exclude_joints: list[str] | None = None + policy_use_bf16: bool | None = None + # Optimizer optimizer_type: str | None = "adam" optimizer_lr: float | None = None @@ -118,6 +131,8 @@ def build_training_command( cmd.extend(["--dataset.root", request.dataset_root]) if request.dataset_episodes: cmd.extend(["--dataset.episodes"] + [str(ep) for ep in request.dataset_episodes]) + if request.dataset_image_transforms_enable: + cmd.extend(["--dataset.image_transforms.enable", "true"]) # Policy cmd.extend(["--policy.type", request.policy_type]) @@ -143,6 +158,29 @@ def build_training_command( if request.policy_push_to_hub and request.policy_repo_id: cmd.extend(["--policy.repo_id", request.policy_repo_id]) + # GR00T-specific policy flags. These options only exist on the groot config, + # so emit them only for --policy.type=groot; sending them to another policy + # would make draccus abort the run. Each is skipped when unset so the groot + # config's own defaults apply. + if request.policy_type == "groot": + if request.policy_base_model_path: + cmd.extend(["--policy.base_model_path", request.policy_base_model_path]) + if request.policy_embodiment_tag: + cmd.extend(["--policy.embodiment_tag", request.policy_embodiment_tag]) + if request.policy_chunk_size is not None: + cmd.extend(["--policy.chunk_size", str(request.policy_chunk_size)]) + if request.policy_n_action_steps is not None: + cmd.extend(["--policy.n_action_steps", str(request.policy_n_action_steps)]) + if request.policy_use_relative_actions is not None: + cmd.extend( + ["--policy.use_relative_actions", "true" if request.policy_use_relative_actions else "false"] + ) + if request.policy_relative_exclude_joints is not None: + # draccus parses list values from a single JSON token, e.g. '["gripper"]'. + cmd.extend(["--policy.relative_exclude_joints", json.dumps(request.policy_relative_exclude_joints)]) + if request.policy_use_bf16 is not None: + cmd.extend(["--policy.use_bf16", "true" if request.policy_use_bf16 else "false"]) + # Logging / checkpointing cmd.extend(["--log_freq", str(request.log_freq)]) cmd.extend(["--save_freq", str(request.save_freq)]) diff --git a/lelab/utils/system.py b/lelab/utils/system.py index 128a0431..a4029303 100644 --- a/lelab/utils/system.py +++ b/lelab/utils/system.py @@ -200,6 +200,9 @@ def handle_install_wandb_extra_status() -> dict[str, Any]: "pi0": ("transformers", "lerobot[pi]"), "pi0_fast": ("transformers", "lerobot[pi]"), "diffusion": ("diffusers", "lerobot[diffusion]"), + # groot pulls in peft (LoRA), diffusers, timm, dm-tree and decord; peft is a + # reliable, cross-platform stand-in for "the lerobot[groot] extra is present". + "groot": ("peft", "lerobot[groot]"), } # One install manager per install target (lerobot[smolvla] / lerobot[pi] / …), diff --git a/tests/test_rollout.py b/tests/test_rollout.py index c0bb1f3d..49200b4c 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -21,6 +21,8 @@ from __future__ import annotations +import json + import pytest @@ -280,11 +282,56 @@ def test_classify_outcome_ok_warns_and_fails() -> None: assert _classify_outcome(1, True, "DeviceNotConnectedError: follower is not connected") == "failed" +def test_rollout_inference_args_uses_rtc_for_relative_action_policies(tmp_path) -> None: + from lelab.rollout import _rollout_inference_args + + policy_dir = tmp_path / "pretrained_model" + policy_dir.mkdir() + (policy_dir / "config.json").write_text( + json.dumps({"type": "groot", "use_relative_actions": True, "n_action_steps": 16}), + encoding="utf-8", + ) + assert _rollout_inference_args(str(policy_dir)) == [ + "--inference.type=rtc", + "--inference.rtc.execution_horizon=16", + "--inference.queue_threshold=0", + ] + + +def test_rollout_inference_args_omits_rtc_for_absolute_policies(tmp_path) -> None: + from lelab.rollout import _rollout_inference_args + + policy_dir = tmp_path / "pretrained_model" + policy_dir.mkdir() + (policy_dir / "config.json").write_text( + json.dumps({"type": "act", "use_relative_actions": False}), + encoding="utf-8", + ) + assert _rollout_inference_args(str(policy_dir)) == [] + + +def test_rollout_inference_args_uses_groot_chunk_default_when_missing(tmp_path) -> None: + from lelab.rollout import _rollout_inference_args + + policy_dir = tmp_path / "pretrained_model" + policy_dir.mkdir() + (policy_dir / "config.json").write_text( + json.dumps({"type": "groot", "use_relative_actions": True}), + encoding="utf-8", + ) + assert _rollout_inference_args(str(policy_dir)) == [ + "--inference.type=rtc", + "--inference.rtc.execution_horizon=16", + "--inference.queue_threshold=0", + ] + + def test_friendly_hint_maps_common_failures() -> None: from lelab.rollout import _friendly_hint assert "gripper" in (_friendly_hint("Motor overload detected") or "").lower() assert "connect" in (_friendly_hint("Failed to connect to the follower") or "").lower() + assert "rtc" in (_friendly_hint("GrootPolicy.select_action does not support relative-action policies") or "").lower() assert _friendly_hint("some unrecognised traceback") is None assert _friendly_hint(None) is None diff --git a/tests/test_train.py b/tests/test_train.py index 58fd0d15..be8a6537 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -175,3 +175,50 @@ def test_local_target_keeps_push_to_hub() -> None: assert _arg_value(cmd, "--policy.push_to_hub") == "true" assert _arg_value(cmd, "--policy.repo_id") == "me/x" assert "--job.target" not in cmd + + +def test_groot_policy_flags_are_emitted_only_for_groot() -> None: + from lelab.train import TrainingRequest, build_training_command + + req = TrainingRequest( + dataset_repo_id="x", + policy_type="groot", + policy_base_model_path="nvidia/GR00T-N1.7-3B", + policy_embodiment_tag="new_embodiment", + policy_chunk_size=16, + policy_n_action_steps=16, + policy_use_relative_actions=True, + policy_relative_exclude_joints=["gripper"], + policy_use_bf16=True, + ) + cmd = build_training_command(req, "/tmp/out") + + assert _arg_value(cmd, "--policy.base_model_path") == "nvidia/GR00T-N1.7-3B" + assert _arg_value(cmd, "--policy.embodiment_tag") == "new_embodiment" + assert _arg_value(cmd, "--policy.chunk_size") == "16" + assert _arg_value(cmd, "--policy.n_action_steps") == "16" + assert _arg_value(cmd, "--policy.use_relative_actions") == "true" + assert _arg_value(cmd, "--policy.relative_exclude_joints") == '["gripper"]' + assert _arg_value(cmd, "--policy.use_bf16") == "true" + + non_groot = build_training_command( + TrainingRequest( + dataset_repo_id="x", + policy_type="act", + policy_base_model_path="nvidia/GR00T-N1.7-3B", + policy_embodiment_tag="new_embodiment", + policy_chunk_size=16, + policy_n_action_steps=16, + policy_use_relative_actions=True, + policy_relative_exclude_joints=["gripper"], + policy_use_bf16=True, + ), + "/tmp/out", + ) + assert "--policy.base_model_path" not in non_groot + assert "--policy.embodiment_tag" not in non_groot + assert "--policy.chunk_size" not in non_groot + assert "--policy.n_action_steps" not in non_groot + assert "--policy.use_relative_actions" not in non_groot + assert "--policy.relative_exclude_joints" not in non_groot + assert "--policy.use_bf16" not in non_groot