close the RECAP critic-policy loop

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