|
|
| import gc |
| import json |
| import os |
| import time |
| from pathlib import Path |
|
|
| import av |
| import numpy as np |
| import pandas as pd |
|
|
| REPO_ROOT = Path(__file__).resolve().parents[2] |
| DEFAULT_OUTPUT = ( |
| REPO_ROOT |
| / "evaluation" |
| / "outputs" |
| / "continuous_6tasks" |
| ) |
| OUTPUT = Path( |
| os.environ.get("CONT_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 |
|
|
|
|
| EPISODES = (50, 150, 250, 350, 450, 550) |
| EXECUTE_STEPS = 5 |
| ACTION_HORIZON = 10 |
| RANDOM_SEED = 20260916 |
|
|
| ARM_DIMS = smoke.ARM_DIMS |
| GRIPPER_DIMS = smoke.GRIPPER_DIMS |
|
|
| 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", |
| ) |
|
|
|
|
| def create_policy(): |
| config = smoke.train_config.get_config( |
| smoke.CONFIG_NAME |
| ) |
|
|
| assets = smoke.replace_dataclass_fields( |
| config.data.assets, |
| assets_dir=str(smoke.CKPT / "assets"), |
| asset_id="real_piper_x_6tasks", |
| ) |
| data = smoke.replace_dataclass_fields( |
| config.data, |
| repo_id=str(smoke.DATASET), |
| assets=assets, |
| ) |
| config = smoke.replace_dataclass_fields( |
| config, |
| data=data, |
| ) |
|
|
| 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}") |
| return config, policy, load_seconds |
|
|
|
|
| def decode_targets_fast(video_path, targets): |
| targets = [float(x) for x in targets] |
| images = [] |
| actual_timestamps = [] |
| errors = [] |
|
|
| with av.open(str(video_path)) as container: |
| stream = container.streams.video[0] |
| time_base = float(stream.time_base) |
|
|
| seek_time = max( |
| 0.0, |
| targets[0] - 2.0 / smoke.FPS, |
| ) |
| container.seek( |
| int(seek_time / time_base), |
| stream=stream, |
| any_frame=False, |
| backward=True, |
| ) |
|
|
| target_index = 0 |
| previous_frame = None |
| previous_timestamp = None |
|
|
| for frame in container.decode(stream): |
| if frame.pts is None: |
| continue |
|
|
| timestamp = float( |
| frame.pts * stream.time_base |
| ) |
|
|
| while ( |
| target_index < len(targets) |
| and timestamp >= targets[target_index] |
| ): |
| target = targets[target_index] |
|
|
| if ( |
| previous_frame is not None |
| and abs(previous_timestamp - target) |
| <= abs(timestamp - target) |
| ): |
| selected_frame = previous_frame |
| selected_timestamp = previous_timestamp |
| else: |
| selected_frame = frame |
| selected_timestamp = timestamp |
|
|
| error = abs(selected_timestamp - target) |
| if error > 2.0 / smoke.FPS: |
| raise RuntimeError( |
| f"Video timestamp error: " |
| f"target={target}, " |
| f"actual={selected_timestamp}, " |
| f"error={error}, " |
| f"path={video_path}" |
| ) |
|
|
| images.append( |
| smoke.chw_uint8( |
| selected_frame.to_ndarray( |
| format="rgb24" |
| ) |
| ) |
| ) |
| actual_timestamps.append( |
| selected_timestamp |
| ) |
| errors.append(error) |
| target_index += 1 |
|
|
| if target_index >= len(targets): |
| break |
|
|
| previous_frame = frame |
| previous_timestamp = timestamp |
|
|
| if len(images) != len(targets): |
| raise RuntimeError( |
| f"Decoded {len(images)}/{len(targets)} " |
| f"targets from {video_path}" |
| ) |
|
|
| return images, actual_timestamps, errors |
|
|
|
|
| def metric(prediction, target, hold, steps): |
| prediction = prediction[:steps] |
| target = target[:steps] |
| hold = hold[:steps] |
|
|
| result = {} |
|
|
| for name, dims in ( |
| ("all", None), |
| ("arm", ARM_DIMS), |
| ("gripper", GRIPPER_DIMS), |
| ): |
| if dims is None: |
| pred_diff = prediction - target |
| hold_diff = hold - target |
| else: |
| pred_diff = ( |
| prediction[:, dims] - target[:, dims] |
| ) |
| hold_diff = hold[:, dims] - target[:, dims] |
|
|
| result[f"model_mae_{name}"] = float( |
| np.mean(np.abs(pred_diff)) |
| ) |
| result[f"hold_mae_{name}"] = float( |
| np.mean(np.abs(hold_diff)) |
| ) |
|
|
| result["model_beats_hold"] = bool( |
| result["model_mae_all"] |
| < result["hold_mae_all"] |
| ) |
| return result |
|
|
|
|
| def add_metrics(row, prefix, values): |
| for key, value in values.items(): |
| row[f"{prefix}_{key}"] = value |
|
|
|
|
| def main(): |
| config, policy, load_seconds = create_policy() |
| rng = np.random.default_rng(RANDOM_SEED) |
|
|
| records = [] |
| overlap_records = [] |
| predictions = [] |
| targets_all = [] |
| executed_rows = [] |
|
|
| global_call_index = 0 |
|
|
| for task_id, episode_index in enumerate(EPISODES): |
| 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) |
| prompt = str(metadata["tasks"][0]) |
|
|
| call_frames = list( |
| range( |
| 0, |
| episode_length - ACTION_HORIZON + 1, |
| EXECUTE_STEPS, |
| ) |
| ) |
|
|
| print( |
| f"===== TASK {task_id} " |
| f"EPISODE {episode_index} =====" |
| ) |
| print(f"prompt={prompt}") |
| print(f"length={episode_length}") |
| print(f"num_calls={len(call_frames)}") |
|
|
| images_by_camera = {} |
| decode_debug = {} |
|
|
| for camera in smoke.CAMERAS: |
| video_path, from_timestamp, _ = ( |
| smoke.video_path_and_range( |
| metadata, |
| camera, |
| ) |
| ) |
| timestamps = [ |
| from_timestamp + frame / smoke.FPS |
| for frame in call_frames |
| ] |
|
|
| images, actual, errors = ( |
| decode_targets_fast( |
| video_path, |
| timestamps, |
| ) |
| ) |
| images_by_camera[camera] = images |
| decode_debug[camera] = { |
| "max_error": float(max(errors)), |
| "mean_error": float( |
| np.mean(errors) |
| ), |
| } |
|
|
| print(f"video_decode={decode_debug}") |
|
|
| policy._model.reset_memory(policy.memory) |
| previous_prediction = None |
| previous_target = None |
| previous_record = None |
|
|
| for call_index, frame_index in enumerate( |
| call_frames |
| ): |
| observation = { |
| "state": states[frame_index].astype( |
| np.float32, |
| copy=False, |
| ), |
| "images": { |
| camera: images_by_camera[camera][ |
| call_index |
| ] |
| for camera in smoke.CAMERAS |
| }, |
| "prompt": prompt, |
| "step": call_index, |
| } |
|
|
| 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) |
|
|
| if prediction.shape != (10, 14): |
| raise RuntimeError( |
| f"Bad prediction shape: " |
| f"{prediction.shape}" |
| ) |
| if not np.isfinite(prediction).all(): |
| raise RuntimeError( |
| f"Non-finite output at " |
| f"episode={episode_index}, " |
| f"frame={frame_index}" |
| ) |
|
|
| exec_metrics = metric( |
| prediction, |
| target, |
| hold, |
| EXECUTE_STEPS, |
| ) |
| full_metrics = metric( |
| prediction, |
| target, |
| hold, |
| ACTION_HORIZON, |
| ) |
|
|
| row = { |
| "global_call_index": global_call_index, |
| "task_id": task_id, |
| "task": prompt, |
| "episode_index": episode_index, |
| "episode_call_index": call_index, |
| "frame_index": frame_index, |
| "memory_phase": call_index % 3, |
| "wall_ms": wall_ms, |
| "policy_infer_ms": float( |
| output["policy_timing"]["infer_ms"] |
| ), |
| } |
| add_metrics(row, "exec5", exec_metrics) |
| add_metrics(row, "full10", full_metrics) |
| records.append(row) |
|
|
| predictions.append(prediction) |
| targets_all.append(target) |
|
|
| for horizon_index in range(EXECUTE_STEPS): |
| for joint_index, joint_name in enumerate( |
| JOINT_NAMES |
| ): |
| pred_value = float( |
| prediction[ |
| horizon_index, |
| joint_index, |
| ] |
| ) |
| gt_value = float( |
| target[ |
| horizon_index, |
| joint_index, |
| ] |
| ) |
| executed_rows.append( |
| { |
| "task_id": task_id, |
| "task": prompt, |
| "episode_index": episode_index, |
| "call_index": call_index, |
| "source_frame": frame_index, |
| "executed_frame": ( |
| frame_index |
| + horizon_index |
| ), |
| "memory_phase": call_index % 3, |
| "horizon_index": horizon_index, |
| "joint_index": joint_index, |
| "joint_name": joint_name, |
| "joint_group": ( |
| "gripper" |
| if joint_index in (6, 13) |
| else "arm" |
| ), |
| "prediction": pred_value, |
| "ground_truth": gt_value, |
| "signed_error": ( |
| pred_value - gt_value |
| ), |
| "absolute_error": abs( |
| pred_value - gt_value |
| ), |
| } |
| ) |
|
|
| if previous_prediction is not None: |
| gt_overlap = ( |
| target[:5] |
| - previous_target[5:10] |
| ) |
| correct_overlap = ( |
| prediction[:5] |
| - previous_prediction[5:10] |
| ) |
| repeat_head = ( |
| prediction[:5] |
| - previous_prediction[:5] |
| ) |
|
|
| overlap_mae = float( |
| np.mean(np.abs(correct_overlap)) |
| ) |
| repeat_mae = float( |
| np.mean(np.abs(repeat_head)) |
| ) |
|
|
| overlap_records.append( |
| { |
| "task_id": task_id, |
| "task": prompt, |
| "episode_index": episode_index, |
| "previous_call_index": ( |
| call_index - 1 |
| ), |
| "current_call_index": call_index, |
| "previous_frame": frame_index - 5, |
| "current_frame": frame_index, |
| "from_memory_phase": ( |
| (call_index - 1) % 3 |
| ), |
| "to_memory_phase": ( |
| call_index % 3 |
| ), |
| "gt_overlap_max_abs": float( |
| np.max(np.abs(gt_overlap)) |
| ), |
| "correct_overlap_mae": overlap_mae, |
| "repeat_head_mae": repeat_mae, |
| "replay_margin": ( |
| repeat_mae - overlap_mae |
| ), |
| "correct_alignment": bool( |
| overlap_mae < repeat_mae |
| ), |
| "gt_head_tail_motion": float( |
| np.mean( |
| np.abs( |
| previous_target[5:10] |
| - previous_target[:5] |
| ) |
| ) |
| ), |
| "boundary_delta_error": float( |
| np.mean( |
| np.abs( |
| ( |
| prediction[0] |
| - previous_prediction[4] |
| ) |
| - ( |
| target[0] |
| - previous_target[4] |
| ) |
| ) |
| ) |
| ), |
| } |
| ) |
|
|
| previous_prediction = prediction |
| previous_target = target |
| previous_record = row |
| global_call_index += 1 |
|
|
| if ( |
| call_index % 25 == 0 |
| or call_index + 1 == len(call_frames) |
| ): |
| print( |
| f"call={call_index:04d}/" |
| f"{len(call_frames) - 1:04d} " |
| f"frame={frame_index:04d} " |
| f"phase={call_index % 3} " |
| f"exec5_mae=" |
| f"{exec_metrics['model_mae_all']:.6f} " |
| f"wall_ms={wall_ms:.1f}", |
| flush=True, |
| ) |
|
|
| del images_by_camera |
| gc.collect() |
|
|
| records_frame = pd.DataFrame(records) |
| overlap_frame = pd.DataFrame(overlap_records) |
| executed_frame = pd.DataFrame(executed_rows) |
|
|
| prediction_array = np.stack(predictions) |
| target_array = np.stack(targets_all) |
|
|
| records_frame.to_csv( |
| OUTPUT / "continuous_records.csv", |
| index=False, |
| ) |
| overlap_frame.to_csv( |
| OUTPUT / "continuous_overlap.csv", |
| index=False, |
| ) |
| executed_frame.to_csv( |
| OUTPUT / "executed_trajectory_long.csv", |
| index=False, |
| ) |
| np.savez_compressed( |
| OUTPUT / "continuous_predictions.npz", |
| predictions=prediction_array, |
| ground_truth=target_array, |
| task_ids=records_frame[ |
| "task_id" |
| ].to_numpy(dtype=np.int64), |
| 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), |
| ) |
|
|
| per_task_rows = [] |
|
|
| for task_id, group in records_frame.groupby( |
| "task_id" |
| ): |
| overlaps = overlap_frame[ |
| overlap_frame["task_id"] == task_id |
| ] |
|
|
| per_task_rows.append( |
| { |
| "task_id": int(task_id), |
| "episode_index": int( |
| group["episode_index"].iloc[0] |
| ), |
| "task": group["task"].iloc[0], |
| "num_calls": int(len(group)), |
| "exec5_model_mae_all": float( |
| group[ |
| "exec5_model_mae_all" |
| ].mean() |
| ), |
| "exec5_hold_mae_all": float( |
| group[ |
| "exec5_hold_mae_all" |
| ].mean() |
| ), |
| "exec5_model_mae_arm": float( |
| group[ |
| "exec5_model_mae_arm" |
| ].mean() |
| ), |
| "exec5_model_mae_gripper": float( |
| group[ |
| "exec5_model_mae_gripper" |
| ].mean() |
| ), |
| "exec5_win_rate": float( |
| group[ |
| "exec5_model_beats_hold" |
| ].mean() |
| ), |
| "full10_model_mae_all": float( |
| group[ |
| "full10_model_mae_all" |
| ].mean() |
| ), |
| "full10_hold_mae_all": float( |
| group[ |
| "full10_hold_mae_all" |
| ].mean() |
| ), |
| "correct_alignment_rate": float( |
| overlaps[ |
| "correct_alignment" |
| ].mean() |
| ), |
| "mean_replay_margin": float( |
| overlaps[ |
| "replay_margin" |
| ].mean() |
| ), |
| "mean_boundary_delta_error": float( |
| overlaps[ |
| "boundary_delta_error" |
| ].mean() |
| ), |
| } |
| ) |
|
|
| per_task = pd.DataFrame(per_task_rows) |
| per_task.to_csv( |
| OUTPUT / "continuous_per_task.csv", |
| index=False, |
| ) |
|
|
| motion_threshold = float( |
| overlap_frame[ |
| "gt_head_tail_motion" |
| ].quantile(0.75) |
| ) |
| high_motion = overlap_frame[ |
| overlap_frame["gt_head_tail_motion"] |
| >= motion_threshold |
| ] |
|
|
| warm = records_frame[ |
| records_frame["global_call_index"] >= 3 |
| ] |
|
|
| summary = { |
| "evaluation_type": ( |
| "continuous teacher-forced full-episode replay" |
| ), |
| "closed_loop": False, |
| "checkpoint": str(smoke.CKPT), |
| "episodes": list(EPISODES), |
| "num_tasks": len(EPISODES), |
| "num_calls": int(len(records_frame)), |
| "num_executed_actions": int( |
| len(records_frame) * EXECUTE_STEPS |
| ), |
| "execute_steps": EXECUTE_STEPS, |
| "action_horizon": ACTION_HORIZON, |
| "policy_load_seconds": load_seconds, |
| "overall": { |
| "exec5_model_mae_all": float( |
| records_frame[ |
| "exec5_model_mae_all" |
| ].mean() |
| ), |
| "exec5_hold_mae_all": float( |
| records_frame[ |
| "exec5_hold_mae_all" |
| ].mean() |
| ), |
| "exec5_model_mae_arm": float( |
| records_frame[ |
| "exec5_model_mae_arm" |
| ].mean() |
| ), |
| "exec5_model_mae_gripper": float( |
| records_frame[ |
| "exec5_model_mae_gripper" |
| ].mean() |
| ), |
| "exec5_win_rate": float( |
| records_frame[ |
| "exec5_model_beats_hold" |
| ].mean() |
| ), |
| "full10_model_mae_all": float( |
| records_frame[ |
| "full10_model_mae_all" |
| ].mean() |
| ), |
| "full10_hold_mae_all": float( |
| records_frame[ |
| "full10_hold_mae_all" |
| ].mean() |
| ), |
| }, |
| "continuity": { |
| "gt_overlap_max_abs": float( |
| overlap_frame[ |
| "gt_overlap_max_abs" |
| ].max() |
| ), |
| "correct_overlap_mae": float( |
| overlap_frame[ |
| "correct_overlap_mae" |
| ].mean() |
| ), |
| "repeat_head_mae": float( |
| overlap_frame[ |
| "repeat_head_mae" |
| ].mean() |
| ), |
| "mean_replay_margin": float( |
| overlap_frame[ |
| "replay_margin" |
| ].mean() |
| ), |
| "correct_alignment_rate": float( |
| overlap_frame[ |
| "correct_alignment" |
| ].mean() |
| ), |
| "mean_boundary_delta_error": float( |
| overlap_frame[ |
| "boundary_delta_error" |
| ].mean() |
| ), |
| }, |
| "high_motion_continuity": { |
| "threshold": motion_threshold, |
| "num_boundaries": int( |
| len(high_motion) |
| ), |
| "mean_replay_margin": float( |
| high_motion[ |
| "replay_margin" |
| ].mean() |
| ), |
| "correct_alignment_rate": float( |
| high_motion[ |
| "correct_alignment" |
| ].mean() |
| ), |
| }, |
| "timing": { |
| "warm_calls": int(len(warm)), |
| "wall_ms_mean": float( |
| warm["wall_ms"].mean() |
| ), |
| "wall_ms_median": float( |
| warm["wall_ms"].median() |
| ), |
| "wall_ms_p95": float( |
| warm["wall_ms"].quantile(0.95) |
| ), |
| "policy_infer_ms_median": float( |
| warm[ |
| "policy_infer_ms" |
| ].median() |
| ), |
| }, |
| } |
|
|
| (OUTPUT / "continuous_summary.json").write_text( |
| json.dumps( |
| summary, |
| indent=2, |
| ensure_ascii=False, |
| ), |
| encoding="utf-8", |
| ) |
|
|
| print("===== CONTINUOUS SUMMARY =====") |
| print( |
| json.dumps( |
| summary, |
| indent=2, |
| ensure_ascii=False, |
| ) |
| ) |
|
|
| print("===== CONTINUOUS PER TASK =====") |
| print(per_task.to_string(index=False)) |
|
|
| print("CONTINUOUS_REPLAY_COMPLETE") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
