Skip to content

feat: held-out validation during training #77

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZouzouWP wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: held-out validation during training #77

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 61 additions & 0 deletions frontend/src/components/training/config/AdvancedCard.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,67 @@ const AdvancedCard: React.FC<ConfigComponentProps> = ({ config, updateConfig })

<Separator className="bg-slate-700" />

{/* Validation */}
<section className="space-y-4">
<SectionHeading>Validation</SectionHeading>
<div className="flex items-center space-x-3">
<Switch
id="validation_enable"
checked={config.dataset_eval_split > 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);
}}
/>
<Label htmlFor="validation_enable" className="text-slate-300">
Hold out a validation split
</Label>
</div>
{config.dataset_eval_split > 0 && (
<>
<div className="grid grid-cols-1 md:grid-cols-2 gap-4">
<div>
<Label htmlFor="dataset_eval_split" className="text-slate-300">
Validation Split (fraction)
</Label>
<NumberInput
id="dataset_eval_split"
integer={false}
step="0.05"
value={config.dataset_eval_split}
onChange={(v) => {
if (v !== undefined) updateConfig('dataset_eval_split', v);
}}
className="bg-slate-900 border-slate-600 text-white rounded-lg"
/>
</div>
<div>
<Label htmlFor="eval_steps" className="text-slate-300">
Evaluate Every N Steps
</Label>
<NumberInput
id="eval_steps"
value={config.eval_steps}
onChange={(v) => {
if (v !== undefined) updateConfig('eval_steps', v);
}}
className="bg-slate-900 border-slate-600 text-white rounded-lg"
/>
</div>
</div>
<p className="text-xs text-slate-500">
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.
</p>
</>
)}
</section>

<Separator className="bg-slate-700" />

{/* Optimizer */}
<section className="space-y-4">
<SectionHeading>Optimizer</SectionHeading>
Expand Down
69 changes: 63 additions & 6 deletions frontend/src/components/training/monitoring/MonitoringStats.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -41,6 +44,7 @@ const MonitoringStats: React.FC<MonitoringStatsProps> = ({
const [lossHistory, setLossHistory] = useState<LossPoint[]>([]);
const [lrHistory, setLrHistory] = useState<LrPoint[]>([]);
const lastStepRef = useRef(0);
const lastEvalLossRef = useRef<number | null>(null);
const { baseUrl, fetchWithHeaders } = useApi();

// Seed the curves from the persisted log on mount (and when the active job
Expand All @@ -53,8 +57,12 @@ const MonitoringStats: React.FC<MonitoringStatsProps> = ({
.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)
Expand All @@ -67,6 +75,10 @@ const MonitoringStats: React.FC<MonitoringStatsProps> = ({
// 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.
Expand All @@ -83,18 +95,41 @@ const MonitoringStats: React.FC<MonitoringStatsProps> = ({
if (step < lastStepRef.current) {
setLossHistory([]);
setLrHistory([]);
lastEvalLossRef.current = null;
}
lastStepRef.current = step;

if (step > 0 && trainingStatus.current_loss != null) {
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) => {
Expand All @@ -103,7 +138,12 @@ const MonitoringStats: React.FC<MonitoringStatsProps> = ({
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
Expand Down Expand Up @@ -160,6 +200,12 @@ const MonitoringStats: React.FC<MonitoringStatsProps> = ({
<span className="text-slate-400 text-sm font-normal">
({trainingStatus.current_loss?.toFixed(4) ?? '—'})
</span>
{trainingStatus.eval_loss != null && (
<span className="text-pink-300 text-sm font-normal">
{' '}
· val {trainingStatus.eval_loss.toFixed(4)}
</span>
)}
</span>
</CardTitle>
</CardHeader>
Expand Down Expand Up @@ -192,15 +238,26 @@ const MonitoringStats: React.FC<MonitoringStatsProps> = ({
borderRadius: 8,
}}
labelStyle={{ color: '#cbd5e1' }}
itemStyle={{ color: '#34d399' }}
formatter={(v: number) => v.toFixed(4)}
/>
<Line
type="monotone"
dataKey="loss"
name="train"
stroke="#34d399"
strokeWidth={2}
dot={false}
connectNulls
isAnimationActive={false}
/>
<Line
type="monotone"
dataKey="eval_loss"
name="val"
stroke="#f472b6"
strokeWidth={2}
dot={{ r: 3, fill: '#f472b6' }}
connectNulls
isAnimationActive={false}
/>
</LineChart>
Expand Down
6 changes: 6 additions & 0 deletions frontend/src/components/training/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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: {
Expand Down
4 changes: 4 additions & 0 deletions frontend/src/lib/jobsApi.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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;
Expand Down
5 changes: 5 additions & 0 deletions frontend/src/pages/Training.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
29 changes: 28 additions & 1 deletion lelab/jobs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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/<entity>/<project>/runs/<id>"
# 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]+")
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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.
Expand All @@ -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)
Expand Down
13 changes: 13 additions & 0 deletions lelab/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@nicolas-rabault nicolas-rabault Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Unbounded. lerobot validates this as 0.0 <= x < 1.0 (configs/default.py:55), but the request model does not, so dataset_eval_split=1.5 is accepted here, the job spawns, and it dies inside make_train_eval_datasets about a minute later with the reason buried in the training log.

Field(ge=0.0, lt=1.0) makes it a 422 on the request, before anything starts.


# Policy configuration
policy_type: str = "act"
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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)])
Comment on lines +128 to +129

@nicolas-rabault nicolas-rabault Aug 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the half of the pairing lerobot does not enforce, and it is the half that fails silently.

configs/train.py:255 rejects eval_steps > 0 with eval_split == 0.0, loudly. The reverse is perfectly legal: dataset_eval_split=0.1 with eval_steps=0 emits --dataset.eval_split 0.1 and nothing else, so make_train_eval_datasets holds out 10% of the episodes, the training set silently shrinks, and no eval ever runs. The user pays for it and gets no curve back, with nothing in the log explaining why.

I hit this by calling build_training_command directly with only dataset_eval_split set, so it is reachable, not hypothetical. The Switch comment in AdvancedCard.tsx already has the right mental model (both or neither); it just needs to live in TrainingRequest as well, as a model_validator, so the invariant holds for every caller of /training/start and not only for the form.


# Policy
cmd.extend(["--policy.type", request.policy_type])
Expand Down Expand Up @@ -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:
Expand Down