|
|
| import json |
| import os |
| import time |
| from pathlib import Path |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import pandas as pd |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
|
|
| REPO_ROOT = Path(__file__).resolve().parents[2] |
| DEFAULT_OUTPUT = ( |
| REPO_ROOT |
| / "evaluation" |
| / "outputs" |
| / "train_replay_6tasks_30clips" |
| ) |
| OUTPUT = Path( |
| os.environ.get("FULL_EVAL_OUTPUT", str(DEFAULT_OUTPUT)) |
| ).expanduser().resolve() |
| OUTPUT.mkdir(parents=True, exist_ok=True) |
|
|
| os.environ.setdefault( |
| "SMOKE_OUTPUT", |
| str(OUTPUT / "_smoke_import"), |
| ) |
|
|
| import offline_replay_smoke as smoke |
|
|
| NUM_EPISODES_PER_TASK = 5 |
| ANCHOR_RATIOS = (0.20, 0.35, 0.50, 0.65, 0.80) |
| EXECUTE_STEPS = 5 |
| ACTION_HORIZON = smoke.ACTION_HORIZON |
| HISTORY_OFFSETS = smoke.HISTORY_OFFSETS |
| RANDOM_SEED = 20260916 |
|
|
| JOINT_NAMES = ( |
| "left_joint_1", |
| "left_joint_2", |
| "left_joint_3", |
| "left_joint_4", |
| "left_joint_5", |
| "left_joint_6", |
| "left_gripper", |
| "right_joint_1", |
| "right_joint_2", |
| "right_joint_3", |
| "right_joint_4", |
| "right_joint_5", |
| "right_joint_6", |
| "right_gripper", |
| ) |
|
|
| ARM_DIMS = smoke.ARM_DIMS |
| GRIPPER_DIMS = smoke.GRIPPER_DIMS |
|
|
|
|
| def load_all_episode_index(): |
| metadata_files = sorted( |
| (smoke.DATASET / "meta" / "episodes").rglob("*.parquet") |
| ) |
| if not metadata_files: |
| raise FileNotFoundError( |
| smoke.DATASET / "meta" / "episodes" |
| ) |
|
|
| tables = [ |
| pq.read_table( |
| path, |
| columns=["episode_index", "tasks", "length"], |
| ) |
| for path in metadata_files |
| ] |
| table = ( |
| pa.concat_tables(tables) |
| if len(tables) > 1 |
| else tables[0] |
| ) |
|
|
| groups = {} |
| for row_index in range(table.num_rows): |
| episode_index = int( |
| table["episode_index"][row_index].as_py() |
| ) |
| tasks = table["tasks"][row_index].as_py() |
| length = int(table["length"][row_index].as_py()) |
|
|
| if not tasks: |
| raise RuntimeError( |
| f"Episode {episode_index} has no task" |
| ) |
|
|
| task = str(tasks[0]) |
| groups.setdefault(task, []).append( |
| { |
| "episode_index": episode_index, |
| "length": length, |
| } |
| ) |
|
|
| ordered_groups = sorted( |
| groups.items(), |
| key=lambda item: min( |
| row["episode_index"] for row in item[1] |
| ), |
| ) |
|
|
| selections = [] |
|
|
| for task_id, (task, episodes) in enumerate(ordered_groups): |
| episodes = sorted( |
| episodes, |
| key=lambda row: row["episode_index"], |
| ) |
|
|
| if len(episodes) < NUM_EPISODES_PER_TASK: |
| raise RuntimeError( |
| f"Task {task!r} only has {len(episodes)} episodes" |
| ) |
|
|
| positions = np.rint( |
| np.linspace( |
| 0, |
| len(episodes) - 1, |
| NUM_EPISODES_PER_TASK, |
| ) |
| ).astype(int) |
|
|
| for sample_id, (position, anchor_ratio) in enumerate( |
| zip(positions, ANCHOR_RATIOS, strict=True) |
| ): |
| selected = episodes[int(position)] |
| selections.append( |
| { |
| "task_id": task_id, |
| "task": task, |
| "sample_id": sample_id, |
| "episode_index": selected["episode_index"], |
| "length": selected["length"], |
| "anchor_ratio": float(anchor_ratio), |
| } |
| ) |
|
|
| if len(ordered_groups) != 6: |
| raise RuntimeError( |
| f"Expected 6 tasks, got {len(ordered_groups)}" |
| ) |
|
|
| if len(selections) != 30: |
| raise RuntimeError( |
| f"Expected 30 selections, got {len(selections)}" |
| ) |
|
|
| return selections |
|
|
|
|
| def create_policy(): |
| config = smoke.train_config.get_config( |
| smoke.CONFIG_NAME |
| ) |
|
|
| data_config = config.data |
| assets_config = getattr(data_config, "assets", None) |
|
|
| if assets_config is not None: |
| assets_config = smoke.replace_dataclass_fields( |
| assets_config, |
| assets_dir=str(smoke.CKPT / "assets"), |
| asset_id="real_piper_x_6tasks", |
| ) |
|
|
| data_config = smoke.replace_dataclass_fields( |
| data_config, |
| repo_id=str(smoke.DATASET), |
| assets=assets_config, |
| ) |
| config = smoke.replace_dataclass_fields( |
| config, |
| data=data_config, |
| ) |
|
|
| print("===== CONFIG =====", flush=True) |
| print(f"config={config.name}", flush=True) |
| print(f"model={config.model}", flush=True) |
| print(f"checkpoint={smoke.CKPT}", flush=True) |
| print(f"dataset={config.data.repo_id}", flush=True) |
| print( |
| f"assets_dir={config.data.assets.assets_dir}", |
| flush=True, |
| ) |
|
|
| start = time.monotonic() |
| policy = smoke.policy_config.create_trained_policy( |
| config, |
| smoke.CKPT, |
| ) |
| load_seconds = time.monotonic() - start |
|
|
| print( |
| f"policy_load_seconds={load_seconds:.3f}", |
| flush=True, |
| ) |
|
|
| return config, policy, load_seconds |
|
|
|
|
| def array_errors(prediction, target, hold, steps): |
| prediction = prediction[:steps] |
| target = target[:steps] |
| hold = hold[:steps] |
|
|
| model_difference = prediction - target |
| hold_difference = hold - target |
|
|
| result = {} |
|
|
| for name, dims in ( |
| ("all", None), |
| ("arm", ARM_DIMS), |
| ("gripper", GRIPPER_DIMS), |
| ): |
| if dims is None: |
| model_selected = model_difference |
| hold_selected = hold_difference |
| else: |
| model_selected = model_difference[:, dims] |
| hold_selected = hold_difference[:, dims] |
|
|
| result[f"model_mae_{name}"] = float( |
| np.mean(np.abs(model_selected)) |
| ) |
| result[f"model_rmse_{name}"] = float( |
| np.sqrt(np.mean(np.square(model_selected))) |
| ) |
| result[f"hold_mae_{name}"] = float( |
| np.mean(np.abs(hold_selected)) |
| ) |
| result[f"hold_rmse_{name}"] = float( |
| np.sqrt(np.mean(np.square(hold_selected))) |
| ) |
|
|
| result["motion_mae_all"] = float( |
| np.mean(np.abs(target - hold)) |
| ) |
| result["model_beats_hold_all"] = bool( |
| result["model_mae_all"] < result["hold_mae_all"] |
| ) |
|
|
| return result |
|
|
|
|
| def add_prefixed_metrics(row, prefix, metrics): |
| for key, value in metrics.items(): |
| row[f"{prefix}_{key}"] = value |
|
|
|
|
| def decode_clip_images(metadata, frame_indices): |
| images_by_frame = { |
| frame_index: {} |
| for frame_index in frame_indices |
| } |
| decode_debug = {} |
|
|
| for camera in smoke.CAMERAS: |
| video_path, from_timestamp, to_timestamp = ( |
| smoke.video_path_and_range(metadata, camera) |
| ) |
|
|
| target_timestamps = [ |
| from_timestamp + frame_index / smoke.FPS |
| for frame_index in frame_indices |
| ] |
|
|
| images, decoded_timestamps, errors = ( |
| smoke.decode_nearest_frames( |
| video_path, |
| target_timestamps, |
| ) |
| ) |
|
|
| for position, frame_index in enumerate(frame_indices): |
| images_by_frame[frame_index][camera] = ( |
| smoke.chw_uint8(images[position]) |
| ) |
|
|
| decode_debug[camera] = { |
| "video_path": str(video_path), |
| "from_timestamp": from_timestamp, |
| "to_timestamp": to_timestamp, |
| "targets": target_timestamps, |
| "decoded": decoded_timestamps, |
| "errors": errors, |
| } |
|
|
| return images_by_frame, decode_debug |
|
|
|
|
| def aggregate_rows(frame): |
| result = {} |
|
|
| for scope in ("exec5", "full10"): |
| for source in ("model", "hold"): |
| for group in ("all", "arm", "gripper"): |
| column = f"{scope}_{source}_mae_{group}" |
| result[column] = float(frame[column].mean()) |
|
|
| result[f"{scope}_win_rate_all"] = float( |
| frame[f"{scope}_model_beats_hold_all"].mean() |
| ) |
| result[f"{scope}_motion_mae_all"] = float( |
| frame[f"{scope}_motion_mae_all"].mean() |
| ) |
| result[f"{scope}_improvement_all"] = ( |
| result[f"{scope}_hold_mae_all"] |
| - result[f"{scope}_model_mae_all"] |
| ) |
|
|
| return result |
|
|
|
|
| def main(): |
| selections = load_all_episode_index() |
|
|
| print("===== EVALUATION SELECTION =====", flush=True) |
| print( |
| pd.DataFrame(selections).to_string(index=False), |
| flush=True, |
| ) |
|
|
| config, policy, load_seconds = create_policy() |
| rng = np.random.default_rng(RANDOM_SEED) |
|
|
| records = [] |
| predictions = [] |
| ground_truths = [] |
| hold_predictions = [] |
| anchor_states = [] |
| trajectory_rows = [] |
| decode_records = [] |
|
|
| total_clips = len(selections) |
| global_call_index = 0 |
|
|
| for clip_index, selection in enumerate(selections, start=1): |
| episode_index = int(selection["episode_index"]) |
| smoke.EPISODE_INDEX = episode_index |
|
|
| metadata = smoke.load_episode_metadata() |
| episode = smoke.load_episode_data(metadata) |
|
|
| states = episode["states"] |
| actions = episode["actions"] |
| episode_length = len(states) |
|
|
| task_list = metadata["tasks"] |
| if not task_list: |
| raise RuntimeError( |
| f"Episode {episode_index} has no task prompt" |
| ) |
| prompt = str(task_list[0]) |
|
|
| if prompt != selection["task"]: |
| raise RuntimeError( |
| f"Task mismatch for episode {episode_index}: " |
| f"{prompt!r} != {selection['task']!r}" |
| ) |
|
|
| lower_anchor = abs(min(HISTORY_OFFSETS)) |
| upper_anchor = episode_length - ACTION_HORIZON |
|
|
| if upper_anchor < lower_anchor: |
| raise RuntimeError( |
| f"Episode {episode_index} too short: " |
| f"{episode_length}" |
| ) |
|
|
| anchor = int( |
| round( |
| lower_anchor |
| + selection["anchor_ratio"] |
| * (upper_anchor - lower_anchor) |
| ) |
| ) |
|
|
| history_indices = [ |
| anchor + offset |
| for offset in HISTORY_OFFSETS |
| ] |
|
|
| images_by_frame, decode_debug = decode_clip_images( |
| metadata, |
| history_indices, |
| ) |
|
|
| decode_records.append( |
| { |
| "episode_index": episode_index, |
| "anchor": anchor, |
| "history_indices": history_indices, |
| "cameras": decode_debug, |
| } |
| ) |
|
|
| policy._model.reset_memory(policy.memory) |
|
|
| print( |
| f"[clip {clip_index:02d}/{total_clips}] " |
| f"task={selection['task_id']} " |
| f"episode={episode_index} " |
| f"length={episode_length} " |
| f"ratio={selection['anchor_ratio']:.2f} " |
| f"frames={history_indices}", |
| flush=True, |
| ) |
|
|
| for memory_phase, frame_index in enumerate(history_indices): |
| observation = { |
| "state": states[frame_index].astype( |
| np.float32, |
| copy=False, |
| ), |
| "images": { |
| camera: images_by_frame[frame_index][camera] |
| for camera in smoke.CAMERAS |
| }, |
| "prompt": prompt, |
| "step": memory_phase, |
| } |
|
|
| noise = rng.standard_normal( |
| ( |
| ACTION_HORIZON, |
| int(config.model.action_dim), |
| ) |
| ).astype(np.float32) |
|
|
| start = time.monotonic() |
| output = policy.infer( |
| observation, |
| noise=noise, |
| ) |
| wall_ms = ( |
| time.monotonic() - start |
| ) * 1000.0 |
|
|
| prediction = np.asarray( |
| output["actions"], |
| dtype=np.float32, |
| ) |
| target = actions[ |
| frame_index : frame_index + ACTION_HORIZON |
| ].astype(np.float32) |
| hold = np.repeat( |
| states[frame_index][None, :], |
| ACTION_HORIZON, |
| axis=0, |
| ).astype(np.float32) |
|
|
| expected_shape = (ACTION_HORIZON, 14) |
| if prediction.shape != expected_shape: |
| raise RuntimeError( |
| f"Prediction shape {prediction.shape}, " |
| f"expected {expected_shape}" |
| ) |
| if target.shape != expected_shape: |
| raise RuntimeError( |
| f"Target shape {target.shape}, " |
| f"expected {expected_shape}" |
| ) |
| if not np.isfinite(prediction).all(): |
| raise RuntimeError( |
| f"Non-finite prediction at " |
| f"episode={episode_index}, frame={frame_index}" |
| ) |
|
|
| exec5_metrics = array_errors( |
| prediction, |
| target, |
| hold, |
| EXECUTE_STEPS, |
| ) |
| full10_metrics = array_errors( |
| prediction, |
| target, |
| hold, |
| ACTION_HORIZON, |
| ) |
|
|
| policy_infer_ms = float( |
| output.get( |
| "policy_timing", |
| {}, |
| ).get("infer_ms", float("nan")) |
| ) |
|
|
| row = { |
| "global_call_index": global_call_index, |
| "clip_index": clip_index - 1, |
| "task_id": int(selection["task_id"]), |
| "task": prompt, |
| "sample_id": int(selection["sample_id"]), |
| "episode_index": episode_index, |
| "episode_length": episode_length, |
| "anchor_ratio": float( |
| selection["anchor_ratio"] |
| ), |
| "anchor_frame": anchor, |
| "memory_phase": memory_phase, |
| "frame_index": frame_index, |
| "policy_infer_ms": policy_infer_ms, |
| "wall_ms": wall_ms, |
| } |
|
|
| add_prefixed_metrics( |
| row, |
| "exec5", |
| exec5_metrics, |
| ) |
| add_prefixed_metrics( |
| row, |
| "full10", |
| full10_metrics, |
| ) |
|
|
| records.append(row) |
| predictions.append(prediction) |
| ground_truths.append(target) |
| hold_predictions.append(hold) |
| anchor_states.append(states[frame_index]) |
|
|
| for horizon_index in range(ACTION_HORIZON): |
| for joint_index, joint_name in enumerate( |
| JOINT_NAMES |
| ): |
| predicted_value = float( |
| prediction[horizon_index, joint_index] |
| ) |
| target_value = float( |
| target[horizon_index, joint_index] |
| ) |
| hold_value = float( |
| hold[horizon_index, joint_index] |
| ) |
| signed_error = ( |
| predicted_value - target_value |
| ) |
|
|
| trajectory_rows.append( |
| { |
| "global_call_index": global_call_index, |
| "task_id": int( |
| selection["task_id"] |
| ), |
| "task": prompt, |
| "episode_index": episode_index, |
| "memory_phase": memory_phase, |
| "frame_index": frame_index, |
| "horizon_index": horizon_index, |
| "is_executed_step": ( |
| horizon_index < EXECUTE_STEPS |
| ), |
| "joint_index": joint_index, |
| "joint_name": joint_name, |
| "joint_group": ( |
| "gripper" |
| if joint_index in (6, 13) |
| else "arm" |
| ), |
| "prediction": predicted_value, |
| "ground_truth": target_value, |
| "hold_baseline": hold_value, |
| "signed_error": signed_error, |
| "absolute_error": abs(signed_error), |
| } |
| ) |
|
|
| print( |
| f" phase={memory_phase} " |
| f"frame={frame_index} " |
| f"exec5_mae={row['exec5_model_mae_all']:.6f} " |
| f"hold={row['exec5_hold_mae_all']:.6f} " |
| f"full10_mae={row['full10_model_mae_all']:.6f} " |
| f"infer_ms={policy_infer_ms:.1f}", |
| flush=True, |
| ) |
|
|
| global_call_index += 1 |
|
|
| records_frame = pd.DataFrame(records) |
| trajectory_frame = pd.DataFrame(trajectory_rows) |
|
|
| prediction_array = np.stack(predictions) |
| target_array = np.stack(ground_truths) |
| hold_array = np.stack(hold_predictions) |
| state_array = np.stack(anchor_states) |
|
|
| records_path = OUTPUT / "records.csv" |
| trajectory_path = OUTPUT / "trajectory_long.csv" |
| prediction_path = OUTPUT / "predictions.npz" |
| decode_path = OUTPUT / "video_decode.json" |
|
|
| records_frame.to_csv(records_path, index=False) |
| trajectory_frame.to_csv(trajectory_path, index=False) |
|
|
| np.savez_compressed( |
| prediction_path, |
| predictions=prediction_array, |
| ground_truth=target_array, |
| hold_baseline=hold_array, |
| states=state_array, |
| episode_indices=records_frame[ |
| "episode_index" |
| ].to_numpy(dtype=np.int64), |
| frame_indices=records_frame[ |
| "frame_index" |
| ].to_numpy(dtype=np.int64), |
| memory_phases=records_frame[ |
| "memory_phase" |
| ].to_numpy(dtype=np.int64), |
| task_ids=records_frame[ |
| "task_id" |
| ].to_numpy(dtype=np.int64), |
| ) |
|
|
| decode_path.write_text( |
| json.dumps( |
| decode_records, |
| indent=2, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| full_history_frame = records_frame[ |
| records_frame["memory_phase"] == 2 |
| ].copy() |
|
|
| overall_all_calls = aggregate_rows(records_frame) |
| overall_full_history = aggregate_rows( |
| full_history_frame |
| ) |
|
|
| per_task_rows = [] |
| for (task_id, task), group in full_history_frame.groupby( |
| ["task_id", "task"], |
| sort=True, |
| ): |
| row = { |
| "task_id": int(task_id), |
| "task": task, |
| "num_clips": int(len(group)), |
| } |
| row.update(aggregate_rows(group)) |
| per_task_rows.append(row) |
|
|
| per_task_frame = pd.DataFrame(per_task_rows) |
| per_task_path = OUTPUT / "per_task.csv" |
| per_task_frame.to_csv(per_task_path, index=False) |
|
|
| per_phase_rows = [] |
| for phase, group in records_frame.groupby( |
| "memory_phase", |
| sort=True, |
| ): |
| row = { |
| "memory_phase": int(phase), |
| "num_calls": int(len(group)), |
| } |
| row.update(aggregate_rows(group)) |
| per_phase_rows.append(row) |
|
|
| per_phase_frame = pd.DataFrame(per_phase_rows) |
| per_phase_path = OUTPUT / "per_memory_phase.csv" |
| per_phase_frame.to_csv(per_phase_path, index=False) |
|
|
| main_indices = full_history_frame.index.to_numpy( |
| dtype=np.int64 |
| ) |
| main_prediction = prediction_array[main_indices] |
| main_target = target_array[main_indices] |
| main_hold = hold_array[main_indices] |
|
|
| per_horizon_rows = [] |
| for horizon_index in range(ACTION_HORIZON): |
| model_difference = ( |
| main_prediction[:, horizon_index] |
| - main_target[:, horizon_index] |
| ) |
| hold_difference = ( |
| main_hold[:, horizon_index] |
| - main_target[:, horizon_index] |
| ) |
|
|
| per_horizon_rows.append( |
| { |
| "horizon_index": horizon_index, |
| "is_executed_step": ( |
| horizon_index < EXECUTE_STEPS |
| ), |
| "model_mae_all": float( |
| np.mean(np.abs(model_difference)) |
| ), |
| "hold_mae_all": float( |
| np.mean(np.abs(hold_difference)) |
| ), |
| "model_mae_arm": float( |
| np.mean( |
| np.abs( |
| model_difference[:, ARM_DIMS] |
| ) |
| ) |
| ), |
| "hold_mae_arm": float( |
| np.mean( |
| np.abs( |
| hold_difference[:, ARM_DIMS] |
| ) |
| ) |
| ), |
| "model_mae_gripper": float( |
| np.mean( |
| np.abs( |
| model_difference[:, GRIPPER_DIMS] |
| ) |
| ) |
| ), |
| "hold_mae_gripper": float( |
| np.mean( |
| np.abs( |
| hold_difference[:, GRIPPER_DIMS] |
| ) |
| ) |
| ), |
| } |
| ) |
|
|
| per_horizon_frame = pd.DataFrame(per_horizon_rows) |
| per_horizon_path = OUTPUT / "per_horizon.csv" |
| per_horizon_frame.to_csv( |
| per_horizon_path, |
| index=False, |
| ) |
|
|
| per_dimension_rows = [] |
| for joint_index, joint_name in enumerate(JOINT_NAMES): |
| model_difference = ( |
| main_prediction[:, :, joint_index] |
| - main_target[:, :, joint_index] |
| ) |
| hold_difference = ( |
| main_hold[:, :, joint_index] |
| - main_target[:, :, joint_index] |
| ) |
|
|
| per_dimension_rows.append( |
| { |
| "joint_index": joint_index, |
| "joint_name": joint_name, |
| "joint_group": ( |
| "gripper" |
| if joint_index in (6, 13) |
| else "arm" |
| ), |
| "model_mae_full10": float( |
| np.mean(np.abs(model_difference)) |
| ), |
| "hold_mae_full10": float( |
| np.mean(np.abs(hold_difference)) |
| ), |
| "model_rmse_full10": float( |
| np.sqrt( |
| np.mean(np.square(model_difference)) |
| ) |
| ), |
| "hold_rmse_full10": float( |
| np.sqrt( |
| np.mean(np.square(hold_difference)) |
| ) |
| ), |
| "model_mae_exec5": float( |
| np.mean( |
| np.abs( |
| model_difference[:, :EXECUTE_STEPS] |
| ) |
| ) |
| ), |
| "hold_mae_exec5": float( |
| np.mean( |
| np.abs( |
| hold_difference[:, :EXECUTE_STEPS] |
| ) |
| ) |
| ), |
| } |
| ) |
|
|
| per_dimension_frame = pd.DataFrame( |
| per_dimension_rows |
| ) |
| per_dimension_path = OUTPUT / "per_dimension.csv" |
| per_dimension_frame.to_csv( |
| per_dimension_path, |
| index=False, |
| ) |
|
|
| motion_threshold = float( |
| full_history_frame[ |
| "full10_motion_mae_all" |
| ].quantile(0.75) |
| ) |
| high_motion_frame = full_history_frame[ |
| full_history_frame[ |
| "full10_motion_mae_all" |
| ] >= motion_threshold |
| ] |
| high_motion_metrics = aggregate_rows( |
| high_motion_frame |
| ) |
|
|
| warm_timing_frame = records_frame[ |
| records_frame["global_call_index"] >= 3 |
| ] |
|
|
| timing_summary = { |
| "cold_first_three_wall_ms": records_frame.iloc[:3][ |
| "wall_ms" |
| ].tolist(), |
| "warm_num_calls": int(len(warm_timing_frame)), |
| "warm_wall_ms_mean": float( |
| warm_timing_frame["wall_ms"].mean() |
| ), |
| "warm_wall_ms_median": float( |
| warm_timing_frame["wall_ms"].median() |
| ), |
| "warm_wall_ms_p95": float( |
| warm_timing_frame["wall_ms"].quantile(0.95) |
| ), |
| "warm_policy_infer_ms_mean": float( |
| warm_timing_frame["policy_infer_ms"].mean() |
| ), |
| "warm_policy_infer_ms_median": float( |
| warm_timing_frame[ |
| "policy_infer_ms" |
| ].median() |
| ), |
| } |
|
|
| summary = { |
| "evaluation_type": ( |
| "teacher-forced offline training-set replay" |
| ), |
| "closed_loop": False, |
| "checkpoint": str(smoke.CKPT), |
| "checkpoint_step": smoke.CHECKPOINT_STEP, |
| "config_name": smoke.CONFIG_NAME, |
| "dataset": str(smoke.DATASET), |
| "num_tasks": 6, |
| "num_clips": int(len(selections)), |
| "calls_per_clip": 3, |
| "num_inference_calls": int(len(records_frame)), |
| "action_horizon": ACTION_HORIZON, |
| "execute_steps": EXECUTE_STEPS, |
| "history_offsets": list(HISTORY_OFFSETS), |
| "random_seed": RANDOM_SEED, |
| "policy_load_seconds": load_seconds, |
| "overall_all_memory_phases": overall_all_calls, |
| "overall_full_history_phase_2": overall_full_history, |
| "high_motion_full_history": { |
| "threshold": motion_threshold, |
| "num_clips": int(len(high_motion_frame)), |
| "metrics": high_motion_metrics, |
| }, |
| "timing": timing_summary, |
| "artifacts": { |
| "records": str(records_path), |
| "trajectory_long": str(trajectory_path), |
| "predictions": str(prediction_path), |
| "video_decode": str(decode_path), |
| "per_task": str(per_task_path), |
| "per_memory_phase": str(per_phase_path), |
| "per_horizon": str(per_horizon_path), |
| "per_dimension": str(per_dimension_path), |
| }, |
| } |
|
|
| summary_path = OUTPUT / "summary.json" |
| summary_path.write_text( |
| json.dumps( |
| summary, |
| indent=2, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| figure, axes = plt.subplots( |
| 2, |
| 2, |
| figsize=(16, 11), |
| ) |
|
|
| x_task = np.arange(len(per_task_frame)) |
| width = 0.38 |
| axes[0, 0].bar( |
| x_task - width / 2, |
| per_task_frame["exec5_model_mae_all"], |
| width, |
| label="StreamPI", |
| ) |
| axes[0, 0].bar( |
| x_task + width / 2, |
| per_task_frame["exec5_hold_mae_all"], |
| width, |
| label="Hold baseline", |
| ) |
| axes[0, 0].set_xticks(x_task) |
| axes[0, 0].set_xticklabels( |
| [ |
| f"Task {task_id}" |
| for task_id in per_task_frame["task_id"] |
| ] |
| ) |
| axes[0, 0].set_title( |
| "Executed 5-step MAE by task (memory phase 2)" |
| ) |
| axes[0, 0].set_ylabel("MAE") |
| axes[0, 0].legend() |
| axes[0, 0].grid(alpha=0.25) |
|
|
| axes[0, 1].plot( |
| per_horizon_frame["horizon_index"], |
| per_horizon_frame["model_mae_all"], |
| marker="o", |
| label="StreamPI", |
| ) |
| axes[0, 1].plot( |
| per_horizon_frame["horizon_index"], |
| per_horizon_frame["hold_mae_all"], |
| marker="o", |
| label="Hold baseline", |
| ) |
| axes[0, 1].axvline( |
| EXECUTE_STEPS - 0.5, |
| color="black", |
| linestyle="--", |
| alpha=0.5, |
| label="Execute boundary", |
| ) |
| axes[0, 1].set_title( |
| "MAE across action horizon (memory phase 2)" |
| ) |
| axes[0, 1].set_xlabel("Horizon index") |
| axes[0, 1].set_ylabel("MAE") |
| axes[0, 1].legend() |
| axes[0, 1].grid(alpha=0.25) |
|
|
| x_joint = np.arange(len(per_dimension_frame)) |
| axes[1, 0].bar( |
| x_joint - width / 2, |
| per_dimension_frame["model_mae_full10"], |
| width, |
| label="StreamPI", |
| ) |
| axes[1, 0].bar( |
| x_joint + width / 2, |
| per_dimension_frame["hold_mae_full10"], |
| width, |
| label="Hold baseline", |
| ) |
| axes[1, 0].set_xticks(x_joint) |
| axes[1, 0].set_xticklabels( |
| per_dimension_frame["joint_name"], |
| rotation=55, |
| ha="right", |
| ) |
| axes[1, 0].set_title( |
| "10-step MAE by action dimension" |
| ) |
| axes[1, 0].set_ylabel("MAE") |
| axes[1, 0].legend() |
| axes[1, 0].grid(alpha=0.25) |
|
|
| axes[1, 1].plot( |
| records_frame["global_call_index"], |
| records_frame["wall_ms"], |
| linewidth=1.0, |
| ) |
| axes[1, 1].axvline( |
| 2.5, |
| color="black", |
| linestyle="--", |
| alpha=0.5, |
| label="Cold compile boundary", |
| ) |
| axes[1, 1].set_title("Inference wall time") |
| axes[1, 1].set_xlabel("Inference call") |
| axes[1, 1].set_ylabel("Milliseconds") |
| axes[1, 1].legend() |
| axes[1, 1].grid(alpha=0.25) |
|
|
| figure.suptitle( |
| "StreamPI 10k ␔ training-set offline replay", |
| fontsize=15, |
| ) |
| figure.tight_layout() |
| plot_path = OUTPUT / "evaluation_overview.png" |
| figure.savefig( |
| plot_path, |
| dpi=180, |
| bbox_inches="tight", |
| ) |
| plt.close(figure) |
|
|
| summary["artifacts"]["plot"] = str(plot_path) |
| summary_path.write_text( |
| json.dumps( |
| summary, |
| indent=2, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| print("===== FULL-HISTORY MAIN METRICS =====") |
| print( |
| json.dumps( |
| overall_full_history, |
| indent=2, |
| ensure_ascii=False, |
| ) |
| ) |
|
|
| print("===== HIGH-MOTION METRICS =====") |
| print( |
| json.dumps( |
| { |
| "threshold": motion_threshold, |
| "num_clips": len(high_motion_frame), |
| "metrics": high_motion_metrics, |
| }, |
| indent=2, |
| ensure_ascii=False, |
| ) |
| ) |
|
|
| print("===== TIMING =====") |
| print( |
| json.dumps( |
| timing_summary, |
| indent=2, |
| ensure_ascii=False, |
| ) |
| ) |
|
|
| print("===== OUTPUTS =====") |
| for name, path in summary["artifacts"].items(): |
| print(f"{name}={path}") |
| print(f"summary={summary_path}") |
| print("FULL_EVAL_COMPLETE") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
