From 2331b1f1a0b1a01d3af02f3e8e7c011805084400 Mon Sep 17 00:00:00 2001 From: ZouzouWP Date: Fri, 17 Jul 2026 17:26:53 +0200 Subject: [PATCH] feat: held-out validation during training MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Until now, training loss was the only signal available during a run — no way to tell a policy was starting to overfit without stopping and evaluating it separately afterward, by which point the ideal checkpoint was already steps behind. ## What changed - New "Validation" section in training's advanced config: a single toggle drives `dataset.eval_split` (fraction of episodes held out from training) and `eval_steps` (evaluation cadence) together, since `lerobot` requires both or neither. - Backend parses `step N: eval_loss=X` lines (emitted by `lerobot_train` whenever `eval_steps > 0`) into a new `eval_loss` field, exposed on both the live training status and the persisted metrics history (so the curve survives a page reload). - The monitoring chart plots `eval_loss` as a second, sparser line alongside the training loss, so overfitting becomes visible while training is still running instead of only in hindsight. ```mermaid flowchart TB A["Held-out episodes\n(e.g. 10%)"] --> B["Periodic evaluation"] C["Training loop"] --> D["Training loss\n(every step)"] B --> E["Validation loss\n(every N steps)"] D --> F["Monitoring chart"] E --> F ``` ## Testing Implemented and code-reviewed against the log-parsing regex and the chart's dedupe/merge logic, but **not yet run end-to-end** with `eval_steps > 0` against a real training job — needs a validation pass to confirm the parsed values match `lerobot_train`'s actual output format before merge. --- .../training/config/AdvancedCard.tsx | 61 ++++++++++++++++ .../training/monitoring/MonitoringStats.tsx | 69 +++++++++++++++++-- frontend/src/components/training/types.ts | 6 ++ frontend/src/lib/jobsApi.ts | 4 ++ frontend/src/pages/Training.tsx | 5 ++ lelab/jobs.py | 29 +++++++- lelab/train.py | 13 ++++ 7 files changed, 180 insertions(+), 7 deletions(-) diff --git a/frontend/src/components/training/config/AdvancedCard.tsx b/frontend/src/components/training/config/AdvancedCard.tsx index 0556303b..ac3f6c91 100644 --- a/frontend/src/components/training/config/AdvancedCard.tsx +++ b/frontend/src/components/training/config/AdvancedCard.tsx @@ -122,6 +122,67 @@ const AdvancedCard: React.FC = ({ config, updateConfig }) + {/* Validation */} +
+ Validation +
+ 0} + onCheckedChange={(checked) => { + // lerobot requires eval_split > 0 whenever eval_steps > 0, + // so the toggle drives both fields together. + updateConfig('dataset_eval_split', checked ? 0.1 : 0); + updateConfig('eval_steps', checked ? 1000 : 0); + }} + /> + +
+ {config.dataset_eval_split > 0 && ( + <> +
+
+ + { + if (v !== undefined) updateConfig('dataset_eval_split', v); + }} + className="bg-slate-900 border-slate-600 text-white rounded-lg" + /> +
+
+ + { + if (v !== undefined) updateConfig('eval_steps', v); + }} + className="bg-slate-900 border-slate-600 text-white rounded-lg" + /> +
+
+

+ Holds out this fraction of episodes from training and computes a + validation loss on them at the given cadence — the eval curve shows + up next to the training loss on the monitoring page. +

+ + )} +
+ + + {/* Optimizer */}
Optimizer diff --git a/frontend/src/components/training/monitoring/MonitoringStats.tsx b/frontend/src/components/training/monitoring/MonitoringStats.tsx index d7a48b4b..6e2cf2fa 100644 --- a/frontend/src/components/training/monitoring/MonitoringStats.tsx +++ b/frontend/src/components/training/monitoring/MonitoringStats.tsx @@ -22,7 +22,10 @@ interface MonitoringStatsProps { interface LossPoint { step: number; - loss: number; + loss?: number; + // Held-out validation loss; sparser cadence than the training loss, so + // most points leave it undefined and the chart line uses connectNulls. + eval_loss?: number; } interface LrPoint { @@ -41,6 +44,7 @@ const MonitoringStats: React.FC = ({ const [lossHistory, setLossHistory] = useState([]); const [lrHistory, setLrHistory] = useState([]); const lastStepRef = useRef(0); + const lastEvalLossRef = useRef(null); const { baseUrl, fetchWithHeaders } = useApi(); // Seed the curves from the persisted log on mount (and when the active job @@ -53,8 +57,12 @@ const MonitoringStats: React.FC = ({ .then((points) => { if (cancelled || points.length === 0) return; const lossSeed: LossPoint[] = points - .filter((p) => p.loss != null) - .map((p) => ({ step: p.step, loss: p.loss as number })) + .filter((p) => p.loss != null || p.eval_loss != null) + .map((p) => ({ + step: p.step, + loss: p.loss ?? undefined, + eval_loss: p.eval_loss ?? undefined, + })) .slice(-HISTORY_CAP); const lrSeed: LrPoint[] = points .filter((p) => p.lr != null) @@ -67,6 +75,10 @@ const MonitoringStats: React.FC = ({ // step-regressed reset in the live-append effect below. const lastSeededStep = points[points.length - 1]?.step ?? 0; lastStepRef.current = lastSeededStep; + // Remember the last seeded eval loss so the live-append effect only + // adds a new eval point when the value actually changes. + const lastEval = [...points].reverse().find((p) => p.eval_loss != null); + lastEvalLossRef.current = lastEval?.eval_loss ?? null; }) .catch(() => { // 404 or transient — fall through; live ticks will populate from empty. @@ -83,6 +95,7 @@ const MonitoringStats: React.FC = ({ if (step < lastStepRef.current) { setLossHistory([]); setLrHistory([]); + lastEvalLossRef.current = null; } lastStepRef.current = step; @@ -90,11 +103,33 @@ const MonitoringStats: React.FC = ({ const loss = trainingStatus.current_loss; setLossHistory((prev) => { const last = prev[prev.length - 1]; - if (last && last.step === step) return prev; + if (last && last.step === step && last.loss != null) return prev; + if (last && last.step === step) { + // Merge onto an eval-only point at the same step. + return [...prev.slice(0, -1), { ...last, loss }]; + } return [...prev, { step, loss }].slice(-HISTORY_CAP); }); } + // The live status only carries the latest eval loss; append a point when + // its value changes (evals run every eval_steps, far sparser than ticks). + if ( + step > 0 && + trainingStatus.eval_loss != null && + trainingStatus.eval_loss !== lastEvalLossRef.current + ) { + const evalLoss = trainingStatus.eval_loss; + lastEvalLossRef.current = evalLoss; + setLossHistory((prev) => { + const last = prev[prev.length - 1]; + if (last && last.step === step) { + return [...prev.slice(0, -1), { ...last, eval_loss: evalLoss }]; + } + return [...prev, { step, eval_loss: evalLoss }].slice(-HISTORY_CAP); + }); + } + if (step > 0 && trainingStatus.current_lr != null) { const lr = trainingStatus.current_lr; setLrHistory((prev) => { @@ -103,7 +138,12 @@ const MonitoringStats: React.FC = ({ return [...prev, { step, lr }].slice(-HISTORY_CAP); }); } - }, [trainingStatus.current_step, trainingStatus.current_loss, trainingStatus.current_lr]); + }, [ + trainingStatus.current_step, + trainingStatus.current_loss, + trainingStatus.current_lr, + trainingStatus.eval_loss, + ]); const progress = getProgressPercentage(); // Until tqdm fires its first progress line, total_steps is 0 — show @@ -160,6 +200,12 @@ const MonitoringStats: React.FC = ({ ({trainingStatus.current_loss?.toFixed(4) ?? '—'}) + {trainingStatus.eval_loss != null && ( + + {' '} + · val {trainingStatus.eval_loss.toFixed(4)} + + )} @@ -192,15 +238,26 @@ const MonitoringStats: React.FC = ({ borderRadius: 8, }} labelStyle={{ color: '#cbd5e1' }} - itemStyle={{ color: '#34d399' }} formatter={(v: number) => v.toFixed(4)} /> + diff --git a/frontend/src/components/training/types.ts b/frontend/src/components/training/types.ts index d0ee5f44..a0bdf9b9 100644 --- a/frontend/src/components/training/types.ts +++ b/frontend/src/components/training/types.ts @@ -13,6 +13,11 @@ export interface TrainingConfig { seed?: number; num_workers: number; + // Validation (held-out split). Both must be > 0 for eval to run; + // the form toggles them together. + dataset_eval_split: number; + eval_steps: number; + // Logging and checkpointing log_freq: number; save_freq: number; @@ -50,6 +55,7 @@ export interface TrainingStatus { current_loss?: number; current_lr?: number; grad_norm?: number; + eval_loss?: number; epoch_time?: number; eta_seconds?: number; available_controls: { diff --git a/frontend/src/lib/jobsApi.ts b/frontend/src/lib/jobsApi.ts index 088b727b..f2ae5f2e 100644 --- a/frontend/src/lib/jobsApi.ts +++ b/frontend/src/lib/jobsApi.ts @@ -9,6 +9,7 @@ export interface TrainingMetrics { current_lr: number | null; grad_norm: number | null; eta_seconds: number | null; + eval_loss: number | null; } export interface LogLine { @@ -21,6 +22,7 @@ export type MetricsHistoryPoint = { loss: number | null; lr: number | null; grad_norm: number | null; + eval_loss: number | null; }; // Mirror of the backend TrainingRequest. The frontend doesn't send all of @@ -32,6 +34,8 @@ export interface TrainingRequest { batch_size: number; seed?: number; num_workers: number; + dataset_eval_split?: number; + eval_steps?: number; log_freq: number; save_freq: number; save_checkpoint: boolean; diff --git a/frontend/src/pages/Training.tsx b/frontend/src/pages/Training.tsx index a361d042..d5fe2b6c 100644 --- a/frontend/src/pages/Training.tsx +++ b/frontend/src/pages/Training.tsx @@ -56,6 +56,7 @@ function jobToStatus(job: JobRecord | null, isStarting: boolean): TrainingStatus current_loss: job.metrics.current_loss ?? undefined, current_lr: job.metrics.current_lr ?? undefined, grad_norm: job.metrics.grad_norm ?? undefined, + eval_loss: job.metrics.eval_loss ?? undefined, eta_seconds: job.metrics.eta_seconds ?? undefined, available_controls: { stop_training: job.state === "running", @@ -76,6 +77,8 @@ function configToRequest(c: TrainingConfig): TrainingRequest { batch_size: c.batch_size, seed: c.seed, num_workers: c.num_workers, + dataset_eval_split: c.dataset_eval_split, + eval_steps: c.eval_steps, log_freq: c.log_freq, save_freq: c.save_freq, save_checkpoint: c.save_checkpoint, @@ -113,6 +116,8 @@ const ConfigurationMode: React.FC = () => { batch_size: 8, seed: 1000, num_workers: 4, + dataset_eval_split: 0, + eval_steps: 0, log_freq: 250, save_freq: 1000, save_checkpoint: true, diff --git a/lelab/jobs.py b/lelab/jobs.py index 10810a16..de4ad08f 100644 --- a/lelab/jobs.py +++ b/lelab/jobs.py @@ -61,6 +61,9 @@ class TrainingMetrics(BaseModel): current_lr: float | None = None grad_norm: float | None = None eta_seconds: float | None = None + # Latest held-out validation loss (only present when the job was started + # with eval_steps > 0 and a dataset eval_split). + eval_loss: float | None = None class LogLine(BaseModel): @@ -119,6 +122,7 @@ class MetricsHistoryPoint(BaseModel): loss: float | None = None lr: float | None = None grad_norm: float | None = None + eval_loss: float | None = None def _pid_alive(pid: int) -> bool: @@ -147,6 +151,10 @@ def wandb_run_url(self) -> str | None: ... # tqdm progress: "Training: 1%|▏ | 125/10000 [02:02<2:36:10, 1.05step/s]" _TQDM_RE = re.compile(r"Training:\s*\d+%[^|]*\|[^|]*\|\s*(\d+)/(\d+)\s*\[(?:[\d:]+)<([\d:]+)") +# Held-out validation loss: "step 40000: eval_loss=1.2345" (lerobot_train, +# emitted every eval_steps when dataset.eval_split > 0). +_EVAL_LOSS_RE = re.compile(r"step\s+(\d+):\s*eval_loss=([-+eE\d.]+)") + # Wandb prints something like "wandb: 🚀 View run at https://wandb.ai///runs/" # when it boots. We capture the first URL of that shape we see. _WANDB_URL_RE = re.compile(r"https://wandb\.ai/[^\s/]+/[^\s/]+/runs/[A-Za-z0-9]+") @@ -204,6 +212,11 @@ def parse_metrics_into(line: str, metrics: TrainingMetrics) -> None: with contextlib.suppress(ValueError): metrics.grad_norm = float(line.split("grdn:")[1].split()[0]) + eval_match = _EVAL_LOSS_RE.search(line) + if eval_match: + with contextlib.suppress(ValueError): + metrics.eval_loss = float(eval_match.group(2)) + except Exception as exc: logger.debug("Error parsing log line %r: %s", line, exc) @@ -991,6 +1004,18 @@ def read_metrics_history(self, job_id: str) -> builtins.list[MetricsHistoryPoint except Exception: continue # skip malformed line, same as read_persisted_logs msg = log_line.message + # Eval lines ("step N: eval_loss=X") use their own format and + # cadence; emit them as eval-only points keyed by their step. + eval_match = _EVAL_LOSS_RE.search(msg) + if eval_match: + with contextlib.suppress(ValueError): + step = int(eval_match.group(1)) + eval_loss = float(eval_match.group(2)) + if points and points[-1].step == step: + points[-1].eval_loss = eval_loss + else: + points.append(MetricsHistoryPoint(step=step, eval_loss=eval_loss)) + continue # Only the log-freq lines carry per-step metric values. # Tqdm lines have a step but no loss/lr — skip them so we # don't emit a flat-line point per tqdm tick. @@ -1006,8 +1031,10 @@ def read_metrics_history(self, job_id: str) -> builtins.list[MetricsHistoryPoint lr=fresh.current_lr, grad_norm=fresh.grad_norm, ) - # Dedupe by step: overwrite on consecutive same-step lines. + # Dedupe by step: overwrite on consecutive same-step lines, + # carrying forward an eval_loss captured for the same step. if points and points[-1].step == point.step: + point.eval_loss = points[-1].eval_loss points[-1] = point else: points.append(point) diff --git a/lelab/train.py b/lelab/train.py index bfdc0226..7fac4d65 100644 --- a/lelab/train.py +++ b/lelab/train.py @@ -35,6 +35,9 @@ class TrainingRequest(BaseModel): dataset_revision: str | None = None dataset_root: str | None = None dataset_episodes: list[int] | None = None + # Fraction of episodes held out for validation (0 disables). lerobot + # requires eval_split > 0 whenever eval_steps > 0. + dataset_eval_split: float = 0.0 # Policy configuration policy_type: str = "act" @@ -71,6 +74,10 @@ class TrainingRequest(BaseModel): eval_n_episodes: int = 10 eval_batch_size: int = 50 eval_use_async_envs: bool = False + # Compute eval loss on the held-out split every N steps (0 disables). + eval_steps: int = 0 + # Cap on total eval samples (0 = use all held-out data). + max_eval_samples: int = 0 # Policy-specific policy_device: str | None = "cuda" @@ -118,6 +125,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_eval_split > 0: + cmd.extend(["--dataset.eval_split", str(request.dataset_eval_split)]) # Policy cmd.extend(["--policy.type", request.policy_type]) @@ -184,6 +193,10 @@ def build_training_command( cmd.extend(["--eval.n_episodes", str(request.eval_n_episodes)]) cmd.extend(["--eval.batch_size", str(request.eval_batch_size)]) cmd.extend(["--eval.use_async_envs", "true" if request.eval_use_async_envs else "false"]) + if request.eval_steps > 0: + cmd.extend(["--eval_steps", str(request.eval_steps)]) + if request.max_eval_samples > 0: + cmd.extend(["--max_eval_samples", str(request.max_eval_samples)]) # Optimizer if request.optimizer_type: