close the RECAP critic-policy loop

src/pi06mini/data.py

@@ -214,7 +214,7 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
         history_steps: int = 8,
         action_horizon: int = 16,
         n_step: int = 4,
-        gamma: float = 0.99,
+        gamma: float = 1.0,
         positive_quantile: float = 0.7,
         seed: int = 7,
         action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None,
@@ -229,6 +229,9 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
         self.seed = int(seed)
         self.state_dim = self.history_steps * 2
         self.action_dim = 2
+        self.n_step = int(n_step)
+        self.gamma = float(gamma)
+        self.positive_quantile = float(positive_quantile)
 
         episodes: list[ExperienceEpisode] = []
         episode_record_indices: list[int] = []
@@ -275,6 +278,12 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
             )
         )
 
+        self.value_histories = torch.stack(
+            [self._episode_value_histories(episode) for episode in self.episodes]
+        )
+        terminal = torch.zeros(len(self.episodes), 1, dtype=self.normalized_returns.dtype)
+        self.value_targets = torch.cat((self.normalized_returns, terminal), dim=1)
+
         if action_statistics is None:
             self.action_normalizer = ActionNormalizer.fit(
                 [episode.history_xy for episode in episodes],
@@ -287,6 +296,58 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
         else:
             self.action_normalizer = ActionNormalizer.from_dict(action_statistics)
 
+    def _episode_value_histories(self, episode: ExperienceEpisode) -> torch.Tensor:
+        """Create observed state windows for every step and the terminal state."""
+
+        path = torch.cat((episode.history_xy, episode.future_xy), dim=0)
+        windows: list[torch.Tensor] = []
+        history_end = len(episode.history_xy)
+        for step in range(len(episode.future_xy) + 1):
+            end = history_end + step
+            start = max(0, end - self.history_steps)
+            window = path[start:end]
+            if len(window) < self.history_steps:
+                pad = window[:1].expand(self.history_steps - len(window), -1)
+                window = torch.cat((pad, window), dim=0)
+            windows.append(window[-self.history_steps :])
+        return torch.stack(windows)
+
+    def relabel_from_value_predictions(
+        self,
+        predicted_values: torch.Tensor,
+        *,
+        force_interventions_positive: bool = True,
+    ) -> None:
+        """Recompute RECAP indicators from the learned critic predictions."""
+
+        predictions = torch.as_tensor(predicted_values, dtype=torch.float32)
+        expected = (len(self.episodes), self.action_horizon + 1)
+        if predictions.shape != expected:
+            raise ValueError(f"predicted_values must have shape {expected}")
+        rewards = stack_episode_field(self.episodes, "rewards")
+        self.advantages = n_step_advantages(
+            rewards,
+            predictions,
+            n_step=self.n_step,
+            gamma=self.gamma,
+        )
+        intervention_masks = stack_episode_field(self.episodes, "intervention_mask").bool()
+        step_interventions = intervention_masks if force_interventions_positive else None
+        self.positive_steps, self.advantage_threshold = quantile_advantage_labels(
+            self.advantages,
+            quantile=self.positive_quantile,
+            intervention_mask=step_interventions,
+        )
+        episode_advantages = self.advantages.mean(dim=-1)
+        episode_interventions = (
+            intervention_masks.any(dim=-1) if force_interventions_positive else None
+        )
+        self.positive_episodes, self.episode_advantage_threshold = quantile_advantage_labels(
+            episode_advantages,
+            quantile=self.positive_quantile,
+            intervention_mask=episode_interventions,
+        )
+
     @property
     def action_statistics(self) -> ActionStatistics:
         return self.action_normalizer.statistics
@@ -334,9 +395,10 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
             "future_xy": episode.future_xy.clone(),
             "rewards": episode.rewards.clone(),
             "returns": self.normalized_returns[episode_index].clone(),
-            # The critic predicts the beginning-of-episode return for this
-            # minimal trajectory-level buffer; the full vector remains
-            # available for diagnostics and advantage recomputation.
+            "value_histories": self.value_histories[episode_index].clone(),
+            "value_targets": self.value_targets[episode_index].clone(),
+            # Scalar aliases are retained for simple diagnostics; value training
+            # and RECAP relabeling use the complete per-state vectors above.
             "return_target": self.normalized_returns[episode_index, 0].clone(),
             "value_target": self.normalized_returns[episode_index, 0].clone(),
             "advantages": self.advantages[episode_index].clone(),

src/pi06mini/infer.py

@@ -15,6 +15,29 @@ from .tokenizer import MiniTokenizer
 from .value import ValueConfig, ValueModel
 
 
+@torch.no_grad()
+def predict_subtask_token(
+    model: Pi06Mini,
+    tokenizer: MiniTokenizer,
+    image: torch.Tensor,
+    history_xy: torch.Tensor,
+    action_horizon: int,
+    device: torch.device,
+) -> int:
+    """Autoregressively produce the semantic token used by low-level control."""
+
+    query = tokenizer.build_subtask_query(history_xy)[None].to(device)
+    zero_actions = torch.zeros(1, action_horizon, model.config.action_dim, device=device)
+    output = model(
+        image,
+        query,
+        zero_actions,
+        torch.zeros(1, device=device),
+        generated_token_start=query.shape[1],
+    )
+    return int(output.text_logits[0, -1].argmax().cpu())
+
+
 @torch.no_grad()
 def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
     checkpoint=torch.load(checkpoint_path,map_location="cpu",weights_only=False)
@@ -35,13 +58,18 @@ def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
     failure=next(dataset[i] for i in range(len(dataset)) if "stop" in dataset[i]["source_type"])
     image=demo["image"][None].to(device)
 
+    predicted_subtask = predict_subtask_token(
+        recap, tokenizer, image, demo["history_xy"], dataset.action_horizon, device
+    )
+
     def condition(model:Pi06Mini, advantage:str|None, x:torch.Tensor,t:torch.Tensor)->torch.Tensor:
-        seq=tokenizer.build_policy_sequence(
-            demo["history_xy"],demo["reasoning"],demo["actions"],advantage,
-            metadata=None,
-        )
-        ids=seq.input_ids[:seq.action_start_index][None].to(device)
-        return model(image,ids,x,t,generated_token_start=seq.generated_token_start).vector_field
+        ids = tokenizer.build_action_condition(
+            demo["history_xy"], predicted_subtask, advantage
+        )[None].to(device)
+        return model(
+            image, ids, x, t,
+            generated_token_start=tokenizer.build_subtask_query(demo["history_xy"]).numel(),
+        ).vector_field
 
     def sample_field(fn,seed:int=2):
         torch.manual_seed(seed)
@@ -77,16 +105,38 @@ def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
 
     value_rows=[]
     for i in range(len(dataset)):
-        s=dataset[i];ids=tokenizer.build_value_input(s["history_xy"],s["source_type"])[None].to(device)
-        prediction=float(value(s["image"][None].to(device),ids).expectation.cpu())
-        value_rows.append({"source_type":s["source_type"],"target":float(s["value_target"]),"prediction":prediction})
-    fig,ax=plt.subplots(figsize=(8,4));x=np.arange(len(value_rows));w=.35
-    ax.bar(x-w/2,[r["target"] for r in value_rows],w,label="return target")
-    ax.bar(x+w/2,[r["prediction"] for r in value_rows],w,label="value prediction")
-    ax.set_xticks(x,[r["source_type"] for r in value_rows],rotation=20,ha="right");ax.set_ylim(-1.05,.05);ax.grid(True,axis="y",alpha=.3);ax.legend();ax.set_title("201-bin distributional value overfit")
+        s=dataset[i]
+        histories=s["value_histories"]
+        ids=torch.stack([tokenizer.build_value_input(history) for history in histories]).to(device)
+        images=s["image"][None].expand(len(histories),-1,-1,-1).to(device)
+        predictions=value(images,ids).expectation.cpu()
+        targets=s["value_targets"]
+        value_rows.append({
+            "source_type":s["source_type"],
+            "targets":targets.tolist(),
+            "predictions":predictions.tolist(),
+            "curve_mse":float(torch.mean((predictions-targets).square())),
+        })
+    fig,axes=plt.subplots(2,2,figsize=(10,7),sharex=True,sharey=True)
+    for ax,row in zip(axes.flat,value_rows):
+        ax.plot(row["targets"],"o-",label="return target")
+        ax.plot(row["predictions"],"x--",label="value prediction")
+        ax.set_title(f'{row["source_type"]} MSE={row["curve_mse"]:.3f}')
+        ax.set_ylim(-1.05,.05);ax.grid(True,alpha=.3)
+    axes[0,0].legend(fontsize=8);fig.supxlabel("episode step");fig.supylabel("normalized value")
     value_path=outdir/"value_comparison.png";fig.tight_layout();fig.savefig(value_path,dpi=170);plt.close(fig)
-    write_json(outdir/"evaluation.json",{"trajectory_metrics":metrics,"values":value_rows})
-    print({"trajectory_metrics":metrics,"values":value_rows})
+    try:
+        predicted_subtask_label = tokenizer.subtasks.decode(predicted_subtask)
+    except ValueError:
+        predicted_subtask_label = f"token_{predicted_subtask}"
+    evaluation = {
+        "predicted_subtask": predicted_subtask_label,
+        "target_subtask": tokenizer.subtasks.classify(demo["reasoning"]),
+        "trajectory_metrics": metrics,
+        "values": value_rows,
+    }
+    write_json(outdir/"evaluation.json",evaluation)
+    print(evaluation)
     return trajectory_path,value_path
 
 

src/pi06mini/tokenizer.py

@@ -191,15 +191,18 @@ class MiniTokenizer:
         parts.append(query)
         return torch.cat(parts)
 
-    def build_value_input(self, history_xy: Tensor, metadata: str | None = None) -> Tensor:
-        """Return image marker plus discretized state, without policy query tokens."""
+    def build_value_input(self, history_xy: Tensor) -> Tensor:
+        """Return observation tokens for the critic without outcome metadata.
+
+        RECAP's value function estimates ``V(o, task)``.  Experience provenance
+        and success/failure labels are training targets, not observable inputs;
+        including them would let the critic solve the toy buffer by leakage.
+        The single-task minimal experiment keeps the task command implicit.
+        """
 
         coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1).cpu()
         specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
-        parts = [specials, coordinate_tokens]
-        metadata_token = self.encode_metadata(metadata)
-        if metadata_token is not None: parts.append(metadata_token.reshape(1))
-        return torch.cat(parts)
+        return torch.cat((specials, coordinate_tokens))
 
     # Backward-compatible name used by the pi0.5 data utilities.
     build_prefix = build_observation
@@ -216,6 +219,28 @@ class MiniTokenizer:
             return None
         raise ValueError("advantage must be positive, negative, omitted, bool, or None")
 
+    def build_subtask_query(self, history_xy: Tensor) -> Tensor:
+        """Observation prefix whose final token predicts the next subtask."""
+
+        return self.build_observation(history_xy, metadata=None)
+
+    def build_action_condition(
+        self,
+        history_xy: Tensor,
+        subtask_token: int | Tensor,
+        advantage: Advantage | bool | None,
+    ) -> Tensor:
+        """Build the inference prefix ``observation -> subtask -> advantage``."""
+
+        pieces = [
+            self.build_observation(history_xy, metadata=None),
+            torch.as_tensor(subtask_token, dtype=torch.long).reshape(1).cpu(),
+        ]
+        advantage_token = self.encode_advantage(advantage)
+        if advantage_token is not None:
+            pieces.append(advantage_token.reshape(1))
+        return torch.cat(pieces)
+
     @staticmethod
     def _orthonormal_dct_matrix(
         horizon: int, *, device: torch.device, dtype: torch.dtype

src/pi06mini/train.py

@@ -63,6 +63,27 @@ def policy_step(
     }
 
 
+@torch.no_grad()
+def predict_dataset_values(
+    value: ValueModel,
+    dataset: Pi06ExperienceDataset,
+    tokenizer: MiniTokenizer,
+    device: torch.device,
+) -> torch.Tensor:
+    """Run the learned critic over every state used for RECAP relabeling."""
+
+    value.eval()
+    predictions: list[torch.Tensor] = []
+    for episode_index in range(len(dataset.episodes)):
+        sample = dataset[episode_index * dataset.augmentation_factor]
+        histories = sample["value_histories"]
+        ids = torch.stack([tokenizer.build_value_input(history) for history in histories])
+        images = sample["image"][None].expand(len(histories), -1, -1, -1)
+        output = value(images.to(device), ids.to(device))
+        predictions.append(output.expectation.cpu())
+    return torch.stack(predictions)
+
+
 def train(config_path: str) -> Path:
     cfg = load_config(config_path)
     root = Path(__file__).resolve().parents[2]
@@ -73,7 +94,8 @@ def train(config_path: str) -> Path:
         manifest, image_size=cfg["data"]["image_size"],
         augmentation_factor=cfg["experience"]["augmentation_factor"],
         action_horizon=cfg["data"]["action_horizon"], n_step=cfg["experience"]["n_step"],
-        gamma=cfg["experience"]["gamma"], positive_quantile=cfg["experience"]["positive_quantile"],
+        gamma=cfg["experience"].get("gamma", 1.0),
+        positive_quantile=cfg["experience"]["positive_quantile"],
     )
     tokenizer = MiniTokenizer()
     base, value = build_models(cfg, tokenizer)
@@ -99,11 +121,31 @@ def train(config_path: str) -> Path:
     value.train()
     for step in range(1, cfg["train"]["value_steps"] + 1):
         sample = dataset[(step - 1) % len(dataset)]
-        ids = tokenizer.build_value_input(sample["history_xy"], sample["source_type"])[None].to(device)
-        output = value(sample["image"][None].to(device), ids, target_values=sample["value_target"][None].to(device))
+        time_index = (step - 1) % (dataset.action_horizon + 1)
+        history_xy = sample["value_histories"][time_index]
+        ids = tokenizer.build_value_input(history_xy)[None].to(device)
+        target = sample["value_targets"][time_index][None].to(device)
+        output = value(sample["image"][None].to(device), ids, target_values=target)
         value_opt.zero_grad(set_to_none=True); output.loss.backward(); value_opt.step()
         if step == 1 or step % cfg["train"]["log_every"] == 0:
-            row={"stage":"value","step":step,"loss":float(output.loss.detach().cpu()),"prediction":float(output.expectation.detach().cpu()),"target":float(sample["value_target"])};history.append(row);print(json.dumps(row))
+            row={"stage":"value","step":step,"loss":float(output.loss.detach().cpu()),"prediction":float(output.expectation.detach().cpu()),"target":float(target.cpu())};history.append(row);print(json.dumps(row))
+
+    critic_values = predict_dataset_values(value, dataset, tokenizer, device)
+    dataset.relabel_from_value_predictions(
+        critic_values,
+        force_interventions_positive=cfg["recap"].get(
+            "force_interventions_positive", True
+        ),
+    )
+    write_json(outdir / "critic_relabeling.json", {
+        "predicted_values": critic_values.tolist(),
+        "advantage_threshold": float(dataset.advantage_threshold),
+        "episode_advantage_threshold": float(dataset.episode_advantage_threshold),
+        "episode_labels": [
+            "positive" if bool(label) else "negative"
+            for label in dataset.positive_episodes
+        ],
+    })
 
     def fit_policy(stage: str, steps: int, conditioned: bool) -> Pi06Mini:
         model, _ = build_models(cfg, tokenizer); model.load_state_dict(base_state); model.to(device).train()