Fidelity-aware reward prototype (world-model gate)

[ShauryaVM]

Opens a concrete wedge on #123 without waiting on simulator integration or the restored `compute_planner_loss` hook (#115).

The issue's design rule is that a reward term only matters if it carries information the reactive policy cannot already infer. Reasoning-band shaping fails that test for the same DPI reason the coupling collapsed to a no-op. This prototype therefore:

1. Converts world-model prediction error into a `[0, 1]` fidelity weight (`exp(-mse / T)`).
2. Gates `base_reward` (and an optional consequence term vs an external preference target) by that fidelity.
3. Leaves the policy-gradient estimator as a thin `soft_advantage_from_reward` helper so it can plug in later behind the planner objective entry point.

No training-loop wiring yet — deliberately a pure helper under `training/losses/` so it can be reviewed and iterated on before anyone hangs RL on it.

How I tested

• `pytest Model/tests/test_fidelity_aware_reward.py`

Related to #123

Made with Cursor

[riita10069]

Please align with the activities of #177.

[riita10069]

Under the contribution standard proposed in #198, I consider this PR AI slop in its current form. #123 proposes falsifiable offline experiments first, but this PR implements an arbitrary reward formula without integration evidence or experimental results. This is not ready for review; please provide the required real-world validation before resubmitting.

[ShauryaVM]

Aligned with #123 v1 (and what #177 should call later — this PR does not import AlpaSim).

The old helper multiplied the whole reward by WM fidelity, which would zero safety/comfort when the world model is junk. The formula is now the issue's one:

R = R_safety + R_progress + 0.5 R_comfort + g * R_wm

R_safety / R_progress / R_comfort are the existing unicycle rollout + comfort_excess_per_sample tensors from RolloutAlignedLoss. g = exp(-mse / T). No reasoning-band term (DPI trap).

Offline ranking, constructed expert (smooth) vs jerky (8 m/s² oscillation). JSON: Model/evaluation/results/fidelity_reward_experiment.json.

Arm	expert R	jerky R	winner	g
handcrafted base	0.00	−0.30	expert	—
faithful WM + misleading consequence	−4.00	−0.30	jerky	1.0
noise WM + same consequence	≈0	−0.30	expert	3e-8

That is the falsifiable check #123 asked for: a noisy world model cannot flip the comfort/progress ranking. A trusted WM can (here the consequence was deliberately adversarial, so you would not ship that preference — the point is the gate).

How I tested: pytest Model/tests/test_fidelity_aware_reward.py (8 passed) and python -m evaluation.fidelity_reward_experiment.

[gcordova10]

Reran it on the branch: 8 tests pass, the JSON reproduces. The gating change is
right — g on the WM term only, and at g = 2.9e-08 the ranking falls back to the
base, which is what #123 asked for.

Four problems in v1_handcrafted_reward, all with prediction set identical to the
expert so only the shaping speaks.

r_safety penalises turning, not leaving a lane. pos is in the ego frame at t=0,
so relu(|pos_y| - 1.75) measures deviation from the initial heading. Sustained
curvature 0.05 costs −1.068 at your 32 steps / 5 m/s, −19.4 at the repo's 64 / 10 m/s.
The whole jerky arm is −0.300, and #168 has spent the week on the model not turning.

The jerky arm has no jerk. 8.0 * sin(2*pi*t/2) is 8.0 * sin(pi*t) — every
integer sample lands on a zero crossing, max|a| = 4.2e-05 on your own _pair(). Its
curvature oscillates at period 3, so max|pos_y| only reaches 0.59 m and r_safety
returns 0.0000 for both arms, which is why the point above never shows in your table.

progress and comfort are imitation terms. -||pos - exp_pos|| is minimised by
matching the expert: ego stopped and expert stopped returns 0.0000, the optimum. And
comfort_excess_per_sample is relu(peak_pred - peak_target), charging only for being
jerkier than the log. Two of the three handcrafted terms are anchored to the expert and
the third is the one that penalises turning — against a docstring that calls v1 the
issue's formula "minus an imitation KL term".

fidelity_temperature=1.0 is a bare default, derived nowhere and overridden by
nothing. It also moves under you: g tends to 1 as the world model improves, so the
gate stops discriminating just when the WM term becomes worth trusting.

The first two are arithmetic. The other two are about what the handcrafted terms are
meant to measure, and until those are settled I do not think the reward is ready to be
trained against. The gate itself is sound; it is what it gates that needs work.

[ShauryaVM]

@gcordova10 The gate stays as-is (g on R_wm only). The four v1_handcrafted_reward bugs are fixed; tests are written so they fail on the old arithmetic. pytest Model/tests/test_fidelity_aware_reward.py: 15 passed.

1. r_safety is cross-track in the intended-path frame, not ego-y. Positions are still integrated in the t=0 ego frame, but safety uses signed distance to the intended polyline (expert path heading along the trajectory, extended as a lane so we do not cap at the last logged point). Same 32-step / 5 m/s / κ=0.05 in-lane turn: max |y| = 6.24 m, old relu(|y|-1.75) cost −1.068, new r_safety = 0. The same curve vs a straight 3.5 m lane still costs < −1. Turning in-lane is free; leaving the lane is not.

2. The jerky arm now has jerk. 8.0 * sin(2π t / 2) with integer step index is sin(π t) and landed on zeros (max|a| = 4.2e-05). The signal is now 8.0 * sin(2π t_sec / 2 + π/4) — time in seconds (the 2 s period as written) plus a phase offset. Re-run: max|a| = 7.90. Curvature still weaves in-lane (max |cte| = 0.99 m < 1.75), so r_safety stays 0 for both arms; comfort is what separates them.

3. Progress and comfort are not imitation. Progress is along-track displacement along the intended path, normalized by the coast horizon (expert = 1.00; a braked “match the stopped log” sample loses to a coast). Comfort is mean jerk / lateral-accel excess vs the physical limits already in RolloutAlignedLoss (4.13 m/s³, 4.89 m/s²), not relu(peak_pred - peak_log). Matching a jerky log no longer zeros the term; expert comfort 0, jerky −8.67. The WM term is still only g * R_wm.

4. T is derived; g saturates. Noise arm is N(0, 10²) on the experiment’s 8-d residual (E[mse]=100). A barely-faithful WM is defined as 10× smaller residual (σ=1 → T = 1.0). g = g_sat * exp(-mse / T) with g_sat = (8+1)/(8+2) = 0.9 (rule of succession: even a perfect reconstruction is not consequence-target fidelity — the misleading preference in this experiment). Barely-faithful (mse=T) sits at 0.331, not 0.9; noise g = 2.6e-08. A good-enough WM no longer snaps to full trust.

Updated JSON (Model/evaluation/results/fidelity_reward_experiment.json):

Arm	expert R	jerky R	winner	g
handcrafted base	1.00	−2.84	expert	—
faithful WM + misleading consequence	−6.68	−2.84	jerky	0.90
noise WM + same consequence	1.00	−2.84	expert	2.6e-08

Same falsifiable check as before: a noisy WM cannot flip the ranking; a trusted one can. The base gap is now comfort (real jerk), not ||pos − exp_pos||.