Align pi0.7 minimal structure with paper
This commit is contained in:
seeker
+250
-32
@@ -1,40 +1,258 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np,torch
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .data import Pi07SteerableDataset
|
||||
from .flow import FlowMatching
|
||||
from .model import Pi07Config,Pi07Mini,HighLevelPolicy,TinyGoalGenerator
|
||||
from .runtime import load_config,choose_device,write_json
|
||||
from .tokenizer import MiniTokenizer,EpisodeMetadata
|
||||
from .model import HighLevelPolicy, Pi07Config, Pi07Mini, TinyGoalGenerator
|
||||
from .rtc import build_rtc_conditioning
|
||||
from .runtime import choose_device, write_json
|
||||
from .tokenizer import EpisodeMetadata, MiniTokenizer
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(checkpoint_path):
|
||||
ck=torch.load(checkpoint_path,map_location='cpu',weights_only=False);cfg=ck['config'];root=Path(__file__).resolve().parents[2];dev=choose_device(cfg['device']);tok=MiniTokenizer();ds=Pi07SteerableDataset(root/'data/overfit/overfit_manifest.jsonl',image_size=cfg['data']['image_size'],history_frames=cfg['data']['history_frames'],action_horizon=cfg['data']['action_horizon'],action_statistics=ck['action_statistics'])
|
||||
pcfg=Pi07Config(**ck['policy_config']);policy=Pi07Mini(pcfg);policy.load_state_dict(ck['policy']);policy.to(dev).eval();high=HighLevelPolicy(pcfg);high.load_state_dict(ck['high_level']);high.to(dev).eval();goal=TinyGoalGenerator(tok.vocab_size);goal.load_state_dict(ck['goal']);goal.to(dev).eval();flow=FlowMatching(cfg['flow']['integration_steps']);s=next(ds[i] for i in range(len(ds)) if ds[i]['policy_kind']=='expert');meta=EpisodeMetadata(**s['token_metadata']);image=s['image'][None].to(dev)
|
||||
prompt,_=tok.encode_prompt('drive safely',None,meta,include_subtask=False);hl=high(image,prompt[None].to(dev),s['history_images'][None].to(dev),s['history_states'][None].to(dev),s['current_state'][None].to(dev));pred_subtask=tok.subtasks.decode(int(hl.logits.argmax(-1).cpu()))
|
||||
full=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta);ids=full.input_ids[:full.action_start_index][None].to(dev)
|
||||
torch.manual_seed(cfg['seed']);g=torch.randn_like(s['subgoal_pixel_values'][None]).to(dev);grid=torch.linspace(0,1,6,device=dev)
|
||||
for i in range(5): g=g+(grid[i+1]-grid[i])*goal(image,ids,g,grid[i].expand(1)).vector_field
|
||||
generated_goal=g.cpu()[0]
|
||||
def sample(name,include_meta=True,use_goal=True,use_history=True,cfg_beta=None,generated=False):
|
||||
seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=include_meta);xids=seq.input_ids[:seq.action_start_index][None].to(dev)
|
||||
goal_pixels=(generated_goal if generated else s['subgoal_pixel_values'])[None].to(dev) if use_goal else None
|
||||
common=dict(pixel_values=image,history_pixel_values=s['history_images'][None].to(dev) if use_history else None,history_states=s['history_states'][None].to(dev) if use_history else None,current_state=s['current_state'][None].to(dev),subgoal_pixel_values=goal_pixels)
|
||||
def field(x,t):
|
||||
if cfg_beta is None:return policy(input_ids=xids,noisy_actions=x,timesteps=t,**common).vector_field
|
||||
un=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=False);uids=un.input_ids[:un.action_start_index][None].to(dev)
|
||||
return policy.metadata_cfg(xids,uids,guidance_scale=cfg_beta,noisy_actions=x,timesteps=t,**common)
|
||||
torch.manual_seed(cfg['evaluation']['fixed_noise_seed']);a=flow.sample(field,(1,ds.action_horizon,2),dev,cfg['flow']['integration_steps'])[0].cpu();return ds.action_normalizer.decode(a,s['history_xy']).numpy()
|
||||
curves={'full_oracle_goal':sample('full'),'generated_goal':sample('generated',generated=True),'no_metadata':sample('no_metadata',False),'no_goal':sample('no_goal',use_goal=False),'no_history':sample('no_history',use_history=False),'cfg_beta2':sample('cfg',cfg_beta=2.0)}
|
||||
target=ds.action_normalizer.decode(s['rtc_target_actions'],s['history_xy']).numpy();valid=s['rtc_valid_mask'].numpy().astype(bool);metrics={k:float(np.mean((v[valid]-target[valid])**2)) for k,v in curves.items()}
|
||||
out=Path(checkpoint_path).parent;fig,ax=plt.subplots(figsize=(9,6));ax.plot(target[:,0],target[:,1],'o-',label='expert target',linewidth=2)
|
||||
for k,v in curves.items():ax.plot(v[:,0],v[:,1],'x--',label=f'{k} ({metrics[k]:.3f})')
|
||||
ax.axis('equal');ax.grid(True,alpha=.3);ax.legend(fontsize=8);ax.set_xlabel('forward x (m)');ax.set_ylabel('left y (m)');ax.set_title('π0.7 context steering and ablations');fig.tight_layout();traj=out/'context_ablation.png';fig.savefig(traj,dpi=170);plt.close(fig)
|
||||
fig,axes=plt.subplots(1,3,figsize=(10,3));axes[0].imshow(image[0].cpu().permute(1,2,0));axes[0].set_title('current observation');axes[1].imshow(s['subgoal_pixel_values'].permute(1,2,0));axes[1].set_title('endpoint goal target');axes[2].imshow(generated_goal.permute(1,2,0).clamp(0,1));axes[2].set_title('generated goal');
|
||||
for a in axes:a.axis('off')
|
||||
fig.tight_layout();goalpath=out/'goal_generation.png';fig.savefig(goalpath,dpi=170);plt.close(fig);write_json(out/'evaluation.json',{'predicted_subtask':pred_subtask,'trajectory_metrics':metrics});print({'predicted_subtask':pred_subtask,'trajectory_metrics':metrics});return traj,goalpath
|
||||
def main():
|
||||
p=argparse.ArgumentParser();p.add_argument('--checkpoint',default='outputs/overfit/checkpoint.pt');evaluate(p.parse_args().checkpoint)
|
||||
if __name__=='__main__':main()
|
||||
def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
|
||||
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
||||
cfg = checkpoint["config"]
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
device = choose_device(cfg["device"])
|
||||
tokenizer = MiniTokenizer()
|
||||
dataset = Pi07SteerableDataset(
|
||||
root / "data/overfit/overfit_manifest.jsonl",
|
||||
image_size=cfg["data"]["image_size"],
|
||||
history_frames=cfg["data"]["history_frames"],
|
||||
action_horizon=cfg["data"]["action_horizon"],
|
||||
action_statistics=checkpoint["action_statistics"],
|
||||
max_rtc_delay_steps=cfg["rtc"]["max_delay_steps"],
|
||||
)
|
||||
policy_config = Pi07Config(**checkpoint["policy_config"])
|
||||
policy = Pi07Mini(policy_config)
|
||||
policy.load_state_dict(checkpoint["policy"])
|
||||
policy.to(device).eval()
|
||||
high_level = HighLevelPolicy(policy_config)
|
||||
high_level.load_state_dict(checkpoint["high_level"])
|
||||
high_level.to(device).eval()
|
||||
goal_generator = TinyGoalGenerator(tokenizer.vocab_size)
|
||||
goal_generator.load_state_dict(checkpoint["goal"])
|
||||
goal_generator.to(device).eval()
|
||||
flow = FlowMatching(cfg["flow"]["integration_steps"])
|
||||
|
||||
sample = next(
|
||||
dataset[index]
|
||||
for index in range(len(dataset))
|
||||
if dataset[index]["policy_kind"] == "expert"
|
||||
)
|
||||
metadata = EpisodeMetadata(**sample["token_metadata"])
|
||||
image = sample["image"][None].to(device)
|
||||
high_level_prompt = tokenizer.encode_high_level_prompt(
|
||||
"drive safely", sample["past_subtasks"], metadata
|
||||
)
|
||||
high_level_output = high_level(
|
||||
image,
|
||||
high_level_prompt[None].to(device),
|
||||
sample["history_images"][None].to(device),
|
||||
sample["history_states"][None].to(device),
|
||||
sample["current_state"][None].to(device),
|
||||
)
|
||||
predicted_subtask = tokenizer.subtasks.decode(
|
||||
int(high_level_output.logits.argmax(-1).cpu())
|
||||
)
|
||||
|
||||
goal_prompt, _ = tokenizer.encode_prompt(
|
||||
"drive safely", predicted_subtask, metadata
|
||||
)
|
||||
goal_noise = torch.randn(
|
||||
sample["future_observation"][None].shape,
|
||||
generator=torch.Generator().manual_seed(cfg["evaluation"]["fixed_noise_seed"]),
|
||||
).to(device)
|
||||
goal_grid = torch.linspace(0.0, 1.0, cfg["evaluation"]["goal_flow_steps"] + 1, device=device)
|
||||
generated_goal = goal_noise
|
||||
for index in range(cfg["evaluation"]["goal_flow_steps"]):
|
||||
time = goal_grid[index].expand(1)
|
||||
generated_goal = generated_goal + (goal_grid[index + 1] - goal_grid[index]) * goal_generator(
|
||||
image, goal_prompt[None].to(device), generated_goal, time
|
||||
).vector_field
|
||||
generated_goal_cpu = generated_goal[0].cpu().clamp(0.0, 1.0)
|
||||
|
||||
def sample_policy(
|
||||
*,
|
||||
include_metadata: bool = True,
|
||||
use_goal: bool = True,
|
||||
use_history: bool = True,
|
||||
guidance_beta: float | None = None,
|
||||
rtc_delay: int = 0,
|
||||
oracle_goal: bool = False,
|
||||
) -> tuple[np.ndarray, torch.Tensor]:
|
||||
conditioned = tokenizer.build_policy_sequence(
|
||||
"drive safely",
|
||||
predicted_subtask,
|
||||
sample["actions"],
|
||||
metadata,
|
||||
include_episode_metadata=include_metadata,
|
||||
)
|
||||
conditioned_ids = conditioned.input_ids[: conditioned.action_start_index][None].to(device)
|
||||
common = {
|
||||
"pixel_values": image,
|
||||
"history_pixel_values": sample["history_images"][None].to(device)
|
||||
if use_history
|
||||
else None,
|
||||
"history_states": sample["history_states"][None].to(device)
|
||||
if use_history
|
||||
else None,
|
||||
"current_state": sample["current_state"][None].to(device),
|
||||
"subgoal_pixel_values": (
|
||||
sample["future_observation"][None].to(device)
|
||||
if oracle_goal
|
||||
else generated_goal
|
||||
)
|
||||
if use_goal
|
||||
else None,
|
||||
}
|
||||
|
||||
if guidance_beta is None:
|
||||
def field(actions: torch.Tensor, times: torch.Tensor) -> torch.Tensor:
|
||||
return policy(
|
||||
input_ids=conditioned_ids,
|
||||
noisy_actions=actions,
|
||||
timesteps=times,
|
||||
**common,
|
||||
).vector_field
|
||||
else:
|
||||
unconditional = tokenizer.build_policy_sequence(
|
||||
"drive safely",
|
||||
predicted_subtask,
|
||||
sample["actions"],
|
||||
metadata,
|
||||
include_episode_metadata=False,
|
||||
)
|
||||
unconditional_ids = unconditional.input_ids[
|
||||
: unconditional.action_start_index
|
||||
][None].to(device)
|
||||
|
||||
def field(actions: torch.Tensor, times: torch.Tensor) -> torch.Tensor:
|
||||
return policy.metadata_cfg(
|
||||
conditioned_ids,
|
||||
unconditional_ids,
|
||||
guidance_scale=guidance_beta,
|
||||
noisy_actions=actions,
|
||||
timesteps=times,
|
||||
**common,
|
||||
)
|
||||
|
||||
rtc = build_rtc_conditioning(
|
||||
sample["actions"],
|
||||
rtc_delay,
|
||||
max_delay_steps=cfg["rtc"]["max_delay_steps"],
|
||||
)
|
||||
torch.manual_seed(cfg["evaluation"]["fixed_noise_seed"])
|
||||
normalized = flow.sample(
|
||||
field,
|
||||
(1, dataset.action_horizon, 2),
|
||||
device,
|
||||
cfg["flow"]["integration_steps"],
|
||||
action_prefix=rtc.action_prefix[None].to(device),
|
||||
delays=torch.tensor([rtc_delay], device=device),
|
||||
)[0].cpu()
|
||||
trajectory = dataset.action_normalizer.decode(
|
||||
normalized, sample["history_xy"]
|
||||
).numpy()
|
||||
return trajectory, normalized
|
||||
|
||||
curves: dict[str, np.ndarray] = {}
|
||||
normalized_predictions: dict[str, torch.Tensor] = {}
|
||||
configurations = {
|
||||
"full_generated_goal": {},
|
||||
"full_oracle_goal": {"oracle_goal": True},
|
||||
"no_metadata": {"include_metadata": False},
|
||||
"no_goal": {"use_goal": False},
|
||||
"no_history": {"use_history": False},
|
||||
"metadata_cfg": {"guidance_beta": cfg["steering"]["cfg_beta"]},
|
||||
}
|
||||
for name, arguments in configurations.items():
|
||||
curves[name], normalized_predictions[name] = sample_policy(**arguments)
|
||||
|
||||
target = sample["future_xy"].numpy()
|
||||
target_actions = sample["actions"]
|
||||
trajectory_metrics = {
|
||||
name: float(np.mean((trajectory - target) ** 2))
|
||||
for name, trajectory in curves.items()
|
||||
}
|
||||
action_metrics = {
|
||||
name: float(torch.mean((actions - target_actions) ** 2))
|
||||
for name, actions in normalized_predictions.items()
|
||||
}
|
||||
|
||||
rtc_metrics: dict[str, dict[str, float]] = {}
|
||||
for delay in cfg["evaluation"]["compare_rtc_delays"]:
|
||||
trajectory, actions = sample_policy(rtc_delay=delay)
|
||||
postfix = torch.arange(dataset.action_horizon) >= delay
|
||||
rtc_metrics[str(delay)] = {
|
||||
"postfix_action_mse": float(
|
||||
torch.mean((actions[postfix] - target_actions[postfix]) ** 2)
|
||||
),
|
||||
"trajectory_mse": float(np.mean((trajectory - target) ** 2)),
|
||||
}
|
||||
|
||||
output_dir = Path(checkpoint_path).parent
|
||||
figure, axis = plt.subplots(figsize=(10, 7))
|
||||
axis.plot(target[:, 0], target[:, 1], "o-", label="expert target", linewidth=2)
|
||||
for name, trajectory in curves.items():
|
||||
axis.plot(
|
||||
trajectory[:, 0],
|
||||
trajectory[:, 1],
|
||||
"x--",
|
||||
label=f"{name} ({trajectory_metrics[name]:.3f})",
|
||||
)
|
||||
axis.axis("equal")
|
||||
axis.grid(True, alpha=0.3)
|
||||
axis.legend(fontsize=8)
|
||||
axis.set_xlabel("forward x (m)")
|
||||
axis.set_ylabel("left y (m)")
|
||||
axis.set_title("pi0.7 predicted-subtask context and RTC evaluation")
|
||||
figure.tight_layout()
|
||||
trajectory_path = output_dir / "context_ablation.png"
|
||||
figure.savefig(trajectory_path, dpi=170)
|
||||
plt.close(figure)
|
||||
|
||||
figure, axes = plt.subplots(1, 3, figsize=(10, 3))
|
||||
axes[0].imshow(image[0].cpu().permute(1, 2, 0))
|
||||
axes[0].set_title("current observation")
|
||||
axes[1].imshow(sample["future_observation"].permute(1, 2, 0))
|
||||
axes[1].set_title("future observation target")
|
||||
axes[2].imshow(generated_goal_cpu.permute(1, 2, 0))
|
||||
axes[2].set_title("generated visual subgoal")
|
||||
for axis in axes:
|
||||
axis.axis("off")
|
||||
figure.tight_layout()
|
||||
goal_path = output_dir / "goal_generation.png"
|
||||
figure.savefig(goal_path, dpi=170)
|
||||
plt.close(figure)
|
||||
|
||||
results = {
|
||||
"target_subtask": tokenizer._classify_subtask(sample["reasoning"]),
|
||||
"predicted_subtask": predicted_subtask,
|
||||
"goal_image_mse": float(
|
||||
torch.mean((generated_goal_cpu - sample["future_observation"]) ** 2)
|
||||
),
|
||||
"trajectory_metrics": trajectory_metrics,
|
||||
"normalized_action_metrics": action_metrics,
|
||||
"rtc_metrics": rtc_metrics,
|
||||
}
|
||||
write_json(output_dir / "evaluation.json", results)
|
||||
print(results)
|
||||
return trajectory_path, goal_path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt")
|
||||
evaluate(parser.parse_args().checkpoint)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user