Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
12 changes: 10 additions & 2 deletions Model/evaluation/faithfulness.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,13 +50,20 @@ def reasoning_intervention_delta(
projection: Optional[Any] = None,
geometry_type: Optional[str] = None,
image_transform: Optional[Any] = None,
**forward_kwargs: Any,
) -> dict[str, float]:
"""Mean trajectory L2 between the reasoning-coupled and bypassed runs.

Args:
model: an ``AutoE2E`` built with ``enable_reasoning=True``.
camera_tiles / map_input / visual_history / egomotion_history: one batch.
projection / geometry_type / image_transform: current geometry ABI.
forward_kwargs: extra forward inputs threaded identically into both runs —
navigation (``route_mask`` / ``map_valid`` / ``route_valid``) so the
delta is read at the checkpoint's real operating point, and a fixed
``initial_noise`` so a stochastic planner uses the SAME noise in the
coupled and bypassed runs (otherwise the delta is noise, not the
intervention).

Returns:
``{"trajectory_l2": float}`` — 0.0 while the coupling gate is untrained.
Expand All @@ -76,7 +83,7 @@ def reasoning_intervention_delta(
model.eval()
restore_buffer = _snapshot_buffer(model)
fwd = dict(projection=projection, geometry_type=geometry_type,
image_transform=image_transform, mode="infer")
image_transform=image_transform, mode="infer", **forward_kwargs)

try:
with torch.no_grad():
Expand Down Expand Up @@ -108,6 +115,7 @@ def horizon_intervention_delta(
projection: Optional[Any] = None,
geometry_type: Optional[str] = None,
image_transform: Optional[Any] = None,
**forward_kwargs: Any,
) -> dict[str, float]:
"""Trajectory delta under a targeted intervention on the horizon tokens.

Expand Down Expand Up @@ -155,7 +163,7 @@ def _perturb(pred):
model.eval()
restore_buffer = _snapshot_buffer(model)
fwd = dict(projection=projection, geometry_type=geometry_type,
image_transform=image_transform, mode="infer")
image_transform=image_transform, mode="infer", **forward_kwargs)

original_forward = head.forward

Expand Down
87 changes: 87 additions & 0 deletions Model/tests/test_faithfulness_forward_kwargs.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
"""Demonstration test for the faithfulness `forward_kwargs` extension.

Proves the two gaps found when trying to wire `reasoning_intervention_delta` into
the route-conditioned, flow-matching eval loop, and that threading `forward_kwargs`
(a fixed `initial_noise` + the navigation inputs) fixes both:

1. Without a fixed `initial_noise`, a stochastic planner draws different noise in
the coupled and bypassed runs, so the "intervention delta" is dominated by
noise rather than by the reasoning intervention (which here is 1e-3).
2. Passing the navigation inputs (`route_mask`) must not raise — the pre-extension
signature had no `**kwargs`, so the eval loop could not thread them at all.

Self-contained: a minimal stochastic stub, no fixture / GPU / network.
"""

from __future__ import annotations

import pytest
import torch
from torch import nn

from evaluation.faithfulness import reasoning_intervention_delta

_T, _D = 8, 2
_BUMP = 1e-3 # the reasoning intervention's effect on the trajectory


class _Reactive(nn.Module):
def __init__(self) -> None:
super().__init__()
self.ReasoningHead = nn.Identity() # present => reasoning coupled


class _StochasticStub(nn.Module):
"""trajectory = reasoning_bump + route_term + noise.

- reasoning_bump: `_BUMP` while `Reactive_E2E.ReasoningHead` is not None (the
intervention sets it to None to bypass);
- noise: uses `initial_noise` if given (fixed across runs), else fresh randn;
- route_term: shifts the operating point when `route_mask` is provided.
"""

def __init__(self) -> None:
super().__init__()
self.Reactive_E2E = _Reactive()

def forward(self, camera, map_input, vis_hist, ego, *,
projection=None, geometry_type=None, image_transform=None,
route_mask=None, initial_noise=None, mode="infer", **_):
b = camera.shape[0]
reasoning_on = self.Reactive_E2E.ReasoningHead is not None
bump = _BUMP if reasoning_on else 0.0
route_term = 0.0 if route_mask is None else 0.5
noise = initial_noise if initial_noise is not None else torch.randn(b, _T, _D)
return torch.zeros(b, _T, _D) + bump + route_term + noise


def _inputs(b: int = 2):
return (torch.randn(b, 7, 3, 256, 256), torch.randn(b, 3, 256, 256),
torch.randn(b, 896), torch.randn(b, 256))


def test_gap_naive_call_is_noise_dominated():
torch.manual_seed(0)
out = reasoning_intervention_delta(_StochasticStub(), *_inputs())
# noise ~ N(0,1) in each run -> delta on the order of 1, not the 1e-3 signal.
assert out["trajectory_l2"] > 0.1


def test_fix_fixed_noise_recovers_the_intervention():
torch.manual_seed(0)
b = 2
fixed = torch.zeros(b, _T, _D) # same noise threaded into both runs
out = reasoning_intervention_delta(_StochasticStub(), *_inputs(b),
initial_noise=fixed)
# noise cancels; only the reasoning bump survives: ||[1e-3, 1e-3]|| = 1e-3*sqrt(2)
assert out["trajectory_l2"] == pytest.approx(_BUMP * 2 ** 0.5, abs=2e-4)


def test_fix_navigation_inputs_are_threaded():
b = 2
fixed = torch.zeros(b, _T, _D)
route = torch.ones(b, 10) # the pre-extension signature would TypeError on this
out = reasoning_intervention_delta(_StochasticStub(), *_inputs(b),
route_mask=route, initial_noise=fixed)
# route is identical in both runs -> it cancels, leaving the intervention only.
assert out["trajectory_l2"] == pytest.approx(_BUMP * 2 ** 0.5, abs=2e-4)
61 changes: 61 additions & 0 deletions Model/tests/test_workflow_training_lifecycle.py
Original file line number Diff line number Diff line change
Expand Up @@ -1688,3 +1688,64 @@ def test_resume_load_keeps_rng_tensors_on_cpu():
keywords = {item.arg: item.value for item in resume_load.keywords}
assert ast.literal_eval(keywords["map_location"]) == "cpu"
assert ast.literal_eval(keywords["weights_only"]) is False


class _ReasoningMetricModel(_MetricModel):
"""`_MetricModel` with a bypassable reasoning head. The trajectory picks up a
fixed bump while ``Reactive_E2E.ReasoningHead`` is present, so the intervention
(which sets the head to None) moves the trajectory by exactly that bump — the
quantity the gate metric must report."""

_BUMP = 1e-3

def __init__(self):
super().__init__()
self.Reactive_E2E = SimpleNamespace(ReasoningHead=object())

def __call__(self, visual, *args, **kwargs):
out = super().__call__(visual, *args, **kwargs)
if self.Reactive_E2E.ReasoningHead is not None:
out = out + self._BUMP
return out


def test_intervention_delta_reported_when_opted_in():
model = _ReasoningMetricModel()
loader = [
(_validation_batch(["sample-b", "sample-a"]), None, "pseudo")
]

metrics = workflows._evaluate_open_loop(
model, loader, torch.device("cpu"), report_intervention=True
)

assert "reasoning_intervention_delta" in metrics
# a fixed initial_noise is threaded into both runs, so the delta reflects only
# the reasoning bump (1e-3 across 128 signals): ||1e-3||_2 = 1e-3 * sqrt(128).
assert metrics["reasoning_intervention_delta"] == pytest.approx(
_ReasoningMetricModel._BUMP * 128 ** 0.5, abs=1e-4
)
assert model.training is True # eval mode restored after the intervention


def test_no_intervention_delta_without_reasoning_head():
model = _MetricModel() # no Reactive_E2E.ReasoningHead
loader = [(_validation_batch(["sample-a"]), None, "pseudo")]

metrics = workflows._evaluate_open_loop(
model, loader, torch.device("cpu"), report_intervention=True
)

assert "reasoning_intervention_delta" not in metrics


def test_no_intervention_delta_by_default():
model = _ReasoningMetricModel()
loader = [(_validation_batch(["sample-a"]), None, "pseudo")]

# report_intervention defaults to False -> no cost, no metric (training path).
metrics = workflows._evaluate_open_loop(
model, loader, torch.device("cpu")
)

assert "reasoning_intervention_delta" not in metrics
31 changes: 31 additions & 0 deletions Platform/pipelines/workflows.py
Original file line number Diff line number Diff line change
Expand Up @@ -870,6 +870,7 @@ def _evaluate_open_loop(
route_swap_counterfactual: bool = False,
include_navigation_records: bool = False,
include_rollout_selector_records: bool = False,
report_intervention: bool = False,
) -> dict:
"""Evaluate one fixed loader and return finite ADE/FDE plus its UID digest."""
import hashlib
Expand Down Expand Up @@ -902,6 +903,11 @@ def _evaluate_open_loop(
route_swap_records: list[dict] = []
rollout_selector_records: list[dict] = []
route_cache: dict[str, dict] = {}
_reactive = getattr(model, "Reactive_E2E", None)
report_delta = report_intervention and (
getattr(_reactive, "ReasoningHead", None) is not None
)
intervention_deltas: list[float] = []
model.eval()
try:
with torch.no_grad():
Expand Down Expand Up @@ -1160,6 +1166,22 @@ def _evaluate_open_loop(
horizon_fde[label].append(
float(horizon_errors[-1])
)
if report_delta and len(intervention_deltas) < 50:
from evaluation.faithfulness import (
reasoning_intervention_delta,
)
intervention_deltas.append(
reasoning_intervention_delta(
model, visual, map_context, vis_hist,
ego_hist, projection=projection,
geometry_type=geometry_type,
route_mask=route_mask, map_valid=map_valid,
route_valid=route_valid,
history_frames=history_frames,
future_frames=future_frames,
initial_noise=initial_noise,
)["trajectory_l2"]
)
if navigation_geometry is not None:
from evaluation.navigation_metrics import (
ROUTE_QUALITY_FIELDS,
Expand Down Expand Up @@ -1313,6 +1335,10 @@ def _evaluate_open_loop(
for label in horizon_steps
},
}
if intervention_deltas:
result["reasoning_intervention_delta"] = float(
np.mean(intervention_deltas)
)
if navigation_geometry is not None:
from evaluation.navigation_metrics import (
summarize_navigation_metrics,
Expand Down Expand Up @@ -6455,6 +6481,7 @@ def _run_evaluation(
device,
training_policy=training_policy,
navigation_geometry=navigation_geometry,
report_intervention=True,
route_swap_counterfactual=(navigation_geometry is not None),
include_navigation_records=(
navigation_records_output is not None
Expand Down Expand Up @@ -6755,6 +6782,10 @@ def _run_evaluation(
for key, value in navigation_metrics.items()
if value is not None
})
if evaluation.get("reasoning_intervention_delta") is not None:
logged_metrics["eval/reasoning_intervention_delta"] = float(
evaluation["reasoning_intervention_delta"]
)
mlflow.log_metrics(logged_metrics)

# Artifacts
Expand Down