franka-gr00t-scripts / john_scripts /groot_policy_inference.py
Add Franka-EEF deploy, reference inference, and lerobot conversion scripts
1c7c303 verified about 2 months ago
60.6 kB
# ruff: noqa

import contextlib
import dataclasses
import datetime
import faulthandler
import os
import signal
import time
from moviepy.editor import ImageSequenceClip
import numpy as np
import pandas as pd
from PIL import Image
from droid.robot_env import RobotEnv
import tqdm
import tyro
import enum
import sys
from collections import defaultdict
import logging
import re
import json
import csv

from gr00t.model.policy_flare import BasePolicy, Gr00tPolicy

from gr00t.data.dataset import LeRobotSingleDataset
from gr00t.experiment.data_config import DATA_CONFIG_MAP


faulthandler.enable()

# DROID data collection frequency -- we slow down execution to match this frequency
DROID_CONTROL_FREQUENCY = 10

class EnvMode(enum.Enum):
    """Supported environments."""
    GROOT = "groot"

@dataclasses.dataclass
class Checkpoint:
    """Load a policy from a trained checkpoint."""

    # Training config name (e.g., "pi0_aloha_sim").
    config: str
    # Checkpoint directory (e.g., "checkpoints/pi0_aloha_sim/exp/10000").
    dir: str
    seed: int = 7
    
@dataclasses.dataclass
class Default:
    """Use the default policy for the given environment."""

@dataclasses.dataclass
class Args:
    # Hardware parameters
    # NOTE : We use only one external cam
    left_camera_id: str = "34022131"  
    wrist_camera_id: str = "10623639"  
    # Second exterior ZED, only used by 3-view data configs (real_droid_joint_3view):
    #   video.exterior_image_1_left <- exterior1_camera_id (34022131)
    #   video.exterior_image_2_left <- left_camera_id     (37149196)
    #   video.wrist_image_left      <- wrist_camera_id    (11845904)
    exterior1_camera_id: str = "34022131"

    # Policy parametersu u
    external_camera: str = "left" # which external camera should be fed to the policy, choose from ["left", "right"]

    # Rollout parameters
    max_timesteps: int = 600
    # How many actions to execute from a predicted action chunk before querying policy server again
    # 16 is usually a good default (equals 1 second of action execution).
    open_loop_horizon: int = 16
    
    env: EnvMode = EnvMode.GROOT 
    
    # If provided, will be used in case the "prompt" key is not present in the data, or if the model doesn't have a default prompt.
    default_prompt: str | None = None
    
    # Specifies how to load the policy. If not provided, the default policy for the environment will be used.
    policy: Checkpoint | Default = dataclasses.field(default_factory=Default)

    # Logging path for results and tsv
    logging_path: str = "results"
    
    wowskin_port_right: str = "/dev/ttyACM0"
    wowskin_port_left: str = "/dev/ttyACM1"
    
    embodiment_tag: str = "new_embodiment"
    data_config: str = "real_droid_cartesian_position"
    denoising_steps: int = 4
    # Override action_head_cfg.num_target_vision_tokens (FLARE future tokens) at load time.
    # Required for the play-teleop 3-view checkpoints: trained with 64, config.json says 32.
    vision_token_num: int | None = None

    # --- Trajectory recording ---
    # Save every episode's full trajectory (obs / state / executed action / predicted
    # chunks) as one .npz + .json per episode, including failures and Ctrl+C cancels.
    save_trajectories: bool = True
    # Square size that recorded camera frames are resized to before storing.
    # Native DROID frames are 1280x720, which costs ~5.5 MB/step (~1.9 GB per
    # 350-step episode) for the two cameras -- far too much to keep many episodes.
    # 256 costs ~0.39 MB/step (~137 MB/episode). Set to 0 to store native resolution.
    traj_image_size: int = 256
    # Compress the .npz (smaller on disk, slower to write).
    traj_compress: bool = True

    # --- World-model (action_sensitive_wm) recording ---
    # How camera frames are stored:
    #   "npz"  -- legacy: every frame as uint8 inside the episode .npz. A 3000-step
    #             3-view episode at 256x256 costs ~1.13 GB on disk and ~1.8 GB of RAM,
    #             because the frames are buffered until the episode ends.
    #   "mp4"  -- one h264 file PER CAMERA, written incrementally during the rollout.
    #             ~3 orders of magnitude smaller and O(1) RAM. This is also the layout
    #             DreamDojo/LeRobot wants (videos/<video_key>/episode_NNNNNN.mp4), so
    #             the converter can stream-copy instead of re-encoding.
    #   "both" -- mp4 plus the legacy npz frames (for debugging only).
    traj_video_mode: str = "mp4"
    # Per-camera mp4 resolution as WxH. Default 256x192 matches the only resolution
    # verified end-to-end through DreamDojo (see the GR-1 converter): 256/192 == 640/480,
    # so DreamDojo's VideoCrop.center_crop becomes a no-op and no double-crop occurs.
    # Native DROID frames are 1280x720 (16:9), so reaching 4:3 needs a center crop of
    # the sides -- done by _crop_resize_for_video below, NOT by squashing the aspect.
    traj_video_size: str = "256x192"
    # fps written into the mp4 headers. Must equal the control frequency (10 Hz) or every
    # downstream timestamp is wrong: the legacy preview video hardcoded 15.
    traj_video_fps: float = float(DROID_CONTROL_FREQUENCY)
    # Side-by-side preview mp4 of all cameras. Useful to eyeball a rollout, useless for
    # training, and expensive: 3x1280x720 @3000 steps = 91 MB/episode. Downscaled by
    # preview_video_scale; set preview_video=False to skip it entirely.
    preview_video: bool = True
    preview_video_scale: float = 0.25

    # --- Action-space perturbation (divergent rollouts for world-model training) ---
    # See gr00t/model/action_head/velocity_noise.py. sigma=0 (the default) is a bit-exact
    # no-op, so this block cannot change ordinary eval runs.
    #
    # sigma is in ACTION units of THIS embodiment (7 absolute joint targets in rad +
    # gripper), so the GR-1 sim numbers do not transfer. CALIBRATE OFFLINE FIRST:
    #   python real_scripts/pnp_scripts/probe_noise_magnitude.py --episode <ep.npz>
    # then start hardware runs at the smallest sigma that clears the R1 control floor.
    # WHERE the noise is injected:
    #   "action"   -- actions += sigma*eps AFTER the Euler update; sigma is the true
    #                 action-space std (rad on the joint targets).
    #   "velocity" -- pred_velocity += sigma*eps BEFORE the Euler update; the Euler step
    #                 scales it by dt (0.25 at 4 denoising steps), so the noise reaching
    #                 the action is dt*sigma*eps. sigma is in flow-velocity units.
    # The same sigma means different things in the two modes -- compare them by measured
    # divergence (probe_noise_magnitude.py), never by sigma.
    noise_mode: str = "action"
    noise_sigma: float = 0.0
    # Which Euler steps of the 4-step sampler get perturbed: "all", or "0-1-2" / "0,1".
    # The sim campaign's variants were steps_0, steps_0-1, steps_0-1-2.
    noise_steps: str = "all"
    # Initial-noise seed. HOLD THIS FIXED across the branches of one tuple so every
    # branch starts the sampler from identical initial noise and only the perturbation
    # differs. None = global RNG (not reproducible, not usable for tuples).
    noise_init_seed: int | None = None
    # Perturbation seed. The per-episode seed is noise_perturb_seed + episode_index, so
    # consecutive episodes in one session are distinct branches of the same tuple.
    noise_perturb_seed: int | None = None
    # Label grouping the episodes collected from one physical start state into a tuple.
    # Set this per manual reset; it lands in the manifest as tuple_id.
    tuple_id: str | None = None

    # --- Hardware safety clamp ---
    # MANDATORY when perturbing: a sigma that is harmless in MuJoCo is a violent joint
    # step on a real arm. Each commanded joint target is clamped to within
    # max_joint_step radians of the MEASURED joint position, so a perturbed chunk can
    # never command a large jump away from where the arm actually is.
    # CALIBRATED FROM REAL CLOSED-LOOP DATA (13k unperturbed + 34k perturbed steps,
    # 2026-09-10), which superseded the offline probe's estimate. The clamp bounds
    #     |commanded joint target - MEASURED joint position|
    # and the arm LAGS the target in closed loop, so the real excursion is ~4x the
    # per-chunk step size the probe measures. The probe's clamp column is therefore a
    # LOWER BOUND, and the 0.15 default it suggested clamped 10% of steps on UNPERTURBED
    # rollouts -- distorting control data, not just perturbed data.
    # Measured excursion percentiles (rad), unperturbed:
    #     p50 0.066   p90 0.151   p95 0.187   p99 0.255   max 0.403
    # Clamp rate by bound (unperturbed / velnoise sigma=0.08):
    #     0.15 -> 10.1% / 9.3%     0.2 -> 3.6% / 3.4%
    #     0.3  ->  0.4% / 0.3%     0.4 -> 0.0% / 0.1%
    # 0.3 leaves the nominal policy essentially untouched (outliers only) while still
    # bounding a runaway command. RAISE ONLY against the FR3's own velocity limits.
    clamp_actions: bool = True
    max_joint_step: float = 0.3
    # Gripper command range (dim 7); commands are clipped, not rate-limited.
    gripper_range: tuple[float, float] = (0.0, 1.0)

    # --- Speed ---
    # torch.compile the VLM backbone and the action-head DiT. Measured on an RTX 5090
    # (checkpoint-20000, 4 denoising steps): model forward 37.3 -> 30.0 ms, full
    # get_action p50 41.8 -> 36.1 ms, at the cost of ~15 s of warmup on the first call.
    # NOTE: with the default open_loop_horizon=16 at 10 Hz the policy is only queried
    # once every 1.6 s, so this saves ~7 ms out of a 1600 ms budget -- it does not make
    # the robot run faster. It also does NOT help the tail: the p95 spikes live in the
    # CPU-side modality transforms (p95 ~160 ms), which torch.compile cannot touch.
    # Off by default; enable only if you lower open_loop_horizon a long way.
    torch_compile: bool = False

    
# Default checkpoints that should be used for each environment.
DEFAULT_CHECKPOINT: dict[EnvMode, Checkpoint] = {
    EnvMode.GROOT : Checkpoint(
        config="groot_droid_arms_only",
        dir="nvidia/GR00T-N1.5-3B",
    ),
    
}

# def create_tactile_env(args : Args):
#     wowskin_right = WowSkin(port=args.wowskin_port_right, compute_baseline=True, baseline_samples=5)
#     time.sleep(0.8) 
#     wowskin_left = WowSkin(port=args.wowskin_port_left, compute_baseline=True, baseline_samples=5)
    
#     return wowskin_right, wowskin_left


def postprocess_action(action):
    new_action = {}
    for k, v in action.items():
        if v.ndim == 1:
            new_action[k] = v[..., None]
        else:
            new_action[k] = v
    return new_action

def get_policy_attr(args : Args, attr):
    return getattr(args.policy, attr, None) if hasattr(args.policy, attr) else None
def none2str(x):
    return str(x) if x is not None else ""

# def create_default_policy(env: EnvMode, *, default_prompt: str | None = None) -> _policy.Policy:
#     """Create a default policy for the given environment."""
#     if checkpoint := DEFAULT_CHECKPOINT.get(env):
#         return _policy_config.create_trained_policy(
#             _config.get_config(checkpoint.config), checkpoint.dir, default_prompt=default_prompt
#         )
#     raise ValueError(f"Unsupported environment mode: {env}")

def create_policy(args: Args) -> BasePolicy:
    """Create a policy from the given arguments."""
    import torch
    data_config_obj = DATA_CONFIG_MAP[args.data_config]
    
    match args.policy:
        case Checkpoint():
            modality_config = data_config_obj.modality_config()
            modality_transform = data_config_obj.transform()
            return Gr00tPolicy(
                model_path = args.policy.dir,
                embodiment_tag = args.embodiment_tag, 
                modality_config= modality_config, 
                modality_transform = modality_transform, 
                denoising_steps = args.denoising_steps, 
                device = "cuda" if torch.cuda.is_available() else "cpu",
                vision_token_num = args.vision_token_num,
            )
            
        # case Default(): #If args.policy is a Default object
        #     return create_default_policy(args.env, default_prompt=args.default_prompt)
        
        
def get_success_rate() -> float:
    success: str | float | None = None
    while not isinstance(success, float):
        success = input(
            "Did the rollout succeed? (enter y for 1.0, n for 0.0, h for 0.5)"
        )
        if success == "y":
            success = 1.0
        elif success == "n":
            success = 0.0
        elif success == "h":
            success = 0.5
        else:
            print(f"Invalid input: {success}")
            continue  # invalid 입력이면 다시 입력받음
        if not (0 <= success <= 1):
            print(f"Success must be a number in [0, 1] but got: {success}")
            success = None  # 범위 벗어나면 다시 입력받음
    return success
    
# We are using Ctrl+C to optionally terminate rollouts early -- however, if we press Ctrl+C while the policy server is
# waiting for a new action chunk, it will raise an exception and the server connection dies.
# This context manager temporarily prevents Ctrl+C and delays it after the server call is complete.
@contextlib.contextmanager
def prevent_keyboard_interrupt():
    """Temporarily prevent keyboard interrupts by delaying them until after the protected code."""
    interrupted = False
    original_handler = signal.getsignal(signal.SIGINT)

    def handler(signum, frame):
        nonlocal interrupted
        interrupted = True

    signal.signal(signal.SIGINT, handler)
    try:
        yield
    finally:
        signal.signal(signal.SIGINT, original_handler)
        if interrupted:
            raise KeyboardInterrupt


def main(args: Args):
    #logging.warning(f"best_of_n: {getattr(args.policy, 'best_of_n', None)}, temperature: {getattr(args.policy, 'temperature', None)}, reward_strategy: {getattr(args.policy, 'reward_strategy', None)}, best_of_n_strategy: {getattr(args.policy, 'best_of_n_strategy', None)}")
    print("Creating the GR00T droid env ...")
    # self._robot = ServerInterface(ip_address=nuc_ip)
    # NOTE : launch_robot is included in the init of RobotEnv
    action_space = "joint_position"  
    env = RobotEnv(action_space=action_space, gripper_action_space="position")
    print("Created the GR00T droid env!")

    # Load Policy Model
    policy = create_policy(args)
    # 3-view configs feed both exterior ZEDs (exterior_image_1_left / exterior_image_2_left) + wrist.
    is_3view = args.data_config in ("real_droid_joint_3view",)

    # Attach the action-space perturbation, if any. With noise_sigma=0 and no init seed
    # this stays None and the sampler is untouched (bit-exact identity).
    velocity_noise = None
    if args.noise_sigma or args.noise_init_seed is not None:
        from gr00t.model.action_head.velocity_noise import VelocityNoise, attach, parse_steps

        velocity_noise = VelocityNoise(
            mode=args.noise_mode,
            sigma=args.noise_sigma,
            steps=parse_steps(args.noise_steps),
            init_seed=args.noise_init_seed,
            perturb_seed=args.noise_perturb_seed,
        )
        head = attach(policy, velocity_noise)
        print(f"[noise] attached to {type(head).__name__}: variant="
              f"{velocity_noise.variant_name()} opts={velocity_noise.effective_options()}")
        if args.noise_sigma and not args.clamp_actions:
            # Refusing rather than warning: this is the one combination that can drive the
            # arm hard, and the plan's recurring root cause is "silent default instead of
            # failure". Pass --clamp-actions explicitly (or --no-clamp-actions in sim).
            raise SystemExit(
                "refusing to run with noise_sigma>0 and clamp_actions=False on hardware. "
                "Perturbed joint targets must be clamped before execution; see "
                "FR3_REALWORLD_AUTOMATION_PLAN.md sec.7.3."
            )
        if args.noise_sigma and args.noise_init_seed is None:
            print("[noise] WARNING: noise_sigma>0 but no --noise-init-seed. Branches will "
                  "differ in their INITIAL sampler noise as well as the perturbation, so "
                  "the episodes are NOT a valid same-start counterfactual tuple.")
    if args.torch_compile:
        import torch

        print("[compile] torch.compile on VLM backbone + action-head DiT (first call warms up ~15 s)")
        policy.model.backbone = torch.compile(policy.model.backbone, dynamic=False)
        policy.model.action_head.model = torch.compile(policy.model.action_head.model, dynamic=False)
    #wowskin_right, wowskin_left = create_tactile_env(args)
    print("Loaded policy model!")

    df = pd.DataFrame(columns=[
        "task", "grasp", "success", "grasp_rate", "success_rate", "num_episodes", "avg_inference_time"
    ])
    task_inference_times = defaultdict(list)

    num_trials = 0
    # Counts every episode STARTED (including cancelled and crashed ones), so artifact
    # names are unique for the whole session. num_trials counts only completed episodes
    # and must never be used to name a file -- see the note where this is incremented.
    attempt_index = 0
    last_instruction = None
    success_list = []
    grasp = 0
    success_cnt = 0
    
    timestamp = datetime.datetime.now().strftime("%y%m%d_%H%M%S")
    policy_folder = f"{timestamp}"

    while True:
        inference_times = []
        all_best_scores_dict = []  # action chunk별 best_score_dict를 저장할 리스트 초기화
        # Free form instruction
        prompt_hint = f" [enter = {args.default_prompt!r}]" if (args.default_prompt and last_instruction is None) else ""
        instruction = input(f"Enter instruction{prompt_hint}: ")
        #instruction = "Pick up the blue cube on the brown box and place it in the golden bowl"
        if last_instruction is not None and instruction == "":
            print("Using last instruction: ", last_instruction)
            instruction = last_instruction
        elif instruction == "" and args.default_prompt:
            print("Using default prompt: ", args.default_prompt)
            instruction = args.default_prompt
            
        task_name = re.sub(r'[^\w\-]', '_', instruction.strip())[:60]
        videos_dir = os.path.join(args.logging_path, args.policy.config, task_name, policy_folder, "videos")
        results_dir = os.path.join(args.logging_path, args.policy.config, task_name, policy_folder, "results")
        episodes_dir = os.path.join(args.logging_path, args.policy.config, task_name, policy_folder, "episodes")
        os.makedirs(videos_dir, exist_ok=True)
        os.makedirs(results_dir, exist_ok=True)
        if args.save_trajectories:
            os.makedirs(episodes_dir, exist_ok=True)
        # tsv 파일명
        tsv_filename = os.path.join(results_dir, f"eval.tsv")
        if not os.path.exists(tsv_filename):
            df.to_csv(tsv_filename, sep='\t', index=False)
        
        # process_tracking.tsv 파일명 (results 폴더 아래)
        process_tracking_tsv = os.path.join(results_dir, "process_tracking.tsv")
        # 파일이 없으면 헤더 작성
        if not os.path.exists(process_tracking_tsv):
            with open(process_tracking_tsv, "w", newline="") as f:
                writer = csv.writer(f, delimiter="\t")
                writer.writerow([
                    "episode_id", "action_chunk_idx", "score_list", "best_score", "episode_score_sum", "episode_action_chunk_sum", "action_chunk", "success"  # success 컬럼 추가
                ])
        
        # Rollout parameters
        actions_from_chunk_completed = 0
        pred_action_chunk = None
        action_chunk_idx = 0
        episode_id = num_trials
        # Monotonic across EVERY episode started in this session, unlike num_trials which
        # is not incremented for cancelled episodes. Artifact names must use this: with
        # num_trials, a cancelled episode and the next attempt both resolve to "ep_001"
        # and the second silently OVERWRITES the first's videos (and its npz), leaving a
        # manifest that describes one episode next to video from another.
        attempt_index += 1

        # Prepare to save video of rollout
        video = []
        bar = tqdm.tqdm(range(args.max_timesteps))
        print("Running rollout... press Ctrl+C to stop early.")
        start_rollout_time = time.time()
        interrupted = False
        txt = []
        output = {}

        # Per-step trajectory buffers (see _save_episode below).
        traj = defaultdict(list)
        episode_start_wall = time.time()

        # Per-camera mp4 writers for this episode. Opened lazily on the first frame, and
        # always closed by _save_episode (which every exit path goes through).
        # Clamp telemetry for this episode (reported in the manifest; see _clamp_action).
        clamp_events, clamp_dims_total, clamp_max_excess = 0, 0, 0.0
        crashed_error = None

        def _clamp_stats():
            """Snapshot the clamp counters. Reads the enclosing scope at call time."""
            n = len(traj.get("action_executed", []))
            return {
                "enabled": bool(args.clamp_actions),
                "max_joint_step_rad": float(args.max_joint_step),
                "clamped_steps": clamp_events,
                # The headline number: if this is not near 0, the perturbation the
                # episode actually executed is smaller than the recorded sigma.
                "clamped_step_rate": (clamp_events / n) if n else 0.0,
                "clamped_dims_total": clamp_dims_total,
                "max_excess_rad": clamp_max_excess,
            }

        # Re-seed the perturbation for this episode. init_seed is applied unchanged, so
        # every episode replays the identical initial sampler noise; the perturb stream
        # advances to seed + episode_index, making consecutive episodes distinct branches.
        if velocity_noise is not None:
            velocity_noise.reset(episode_index=episode_id)
            print(f"[noise] episode {episode_id}: {velocity_noise.effective_options()}")

        vid_writer = None
        if args.save_trajectories and args.traj_video_mode in ("mp4", "both"):
            vid_writer = _PerCameraVideoWriter(
                out_dir=os.path.join(episodes_dir, "videos"),
                stem=f"ep_{attempt_index:03d}",
                size=args.traj_video_size,
                fps=args.traj_video_fps,
            )

        for t_step in bar:
            start_time = time.time()
            try:
                # Get the current observation   
                curr_obs = _extract_observation(
                    args,
                    env.get_observation(),
                )
               
                
                exterior_image = curr_obs["left_image"]
                wrist_image = curr_obs["wrist_image"]
                exterior1_image = curr_obs["exterior1_image"]  # None unless the 2nd exterior ZED is connected

                if is_3view:
                    if exterior1_image is None:
                        raise RuntimeError(
                            f"3-view data config {args.data_config} needs camera {args.exterior1_camera_id} "
                            f"(video.exterior_image_1_left) but it is not in the observation"
                        )
                    concat_image = np.concatenate([exterior1_image, exterior_image, wrist_image], axis=1)
                else:
                    concat_image = np.concatenate([exterior_image, wrist_image], axis=1)
                # Downscale BEFORE buffering: the concat frame is 3x1280x720x3 = 8.3 MB, so
                # 3000 steps of full-res preview would hold ~25 GB of RAM until episode end.
                if args.preview_video:
                    video.append(_downscale_preview(concat_image, args.preview_video_scale))

                if args.save_trajectories:
                    # Cameras -> per-camera mp4, keyed by the GR00T model video key so the
                    # output tree is already in DreamDojo's videos/<video_key>/ shape.
                    if vid_writer is not None:
                        if is_3view:
                            vid_writer.append("video.exterior_image_1_left", exterior1_image)
                            vid_writer.append("video.exterior_image_2_left", exterior_image)
                        else:
                            vid_writer.append("video.exterior_image_1_left", exterior_image)
                        vid_writer.append("video.wrist_image_left", wrist_image)
                    # Legacy in-npz frames: only when explicitly asked for. This is what
                    # made episodes 1.1 GB; "mp4" mode stores no frames in the npz.
                    if args.traj_video_mode in ("npz", "both"):
                        traj["exterior_image"].append(_resize_for_traj(exterior_image, args.traj_image_size))
                        traj["wrist_image"].append(_resize_for_traj(wrist_image, args.traj_image_size))
                        if exterior1_image is not None:
                            traj["exterior1_image"].append(_resize_for_traj(exterior1_image, args.traj_image_size))
                    traj["joint_pos_abs"].append(np.asarray(curr_obs["joint_position"], dtype=np.float32))
                    traj["gripper_close"].append(np.asarray(curr_obs["gripper_position"], dtype=np.float32))
                    traj["eef_position"].append(np.asarray(curr_obs["eef_position"], dtype=np.float32))
                    traj["eef_rotation"].append(np.asarray(curr_obs["eef_rotation"], dtype=np.float32))
                    traj["t_step"].append(t_step)
                    traj["time_rel"].append(time.time() - episode_start_wall)
                
                # tactile : [2, 15]
                #tactile_right_info = wowskin_right.read(as_dict=False) 
                #tactile_left_info = wowskin_left.read(as_dict=False) 
                #tactile_left_flat = tactile_left_info.flatten()
                #tactile_right_flat = tactile_right_info.flatten()
                #tactile = np.stack([tactile_left_flat, tactile_right_flat])
                
                if actions_from_chunk_completed == 0 or actions_from_chunk_completed >= args.open_loop_horizon:
                    actions_from_chunk_completed = 0

                    # We resize images on the robot laptop to minimize the amount of data sent to the policy server
                    # and improve latency.

                    # lerobot format for GR00T:
                    import cv2

                    # FOR DEBUG
                    resized_exterior_image = cv2.resize(curr_obs['left_image'], (256, 256), interpolation=cv2.INTER_LINEAR)
                    resized_wrist_image =   cv2.resize(curr_obs['wrist_image'], (256, 256), interpolation=cv2.INTER_LINEAR)
                    mask =  np.ones((256, 256), dtype=np.uint8) * 255
                    mask_3ch = np.stack([mask]*3, axis=-1) 


                    # for key in curr_obs:
                    #     print(key)
                    #     print(curr_obs[key].shape)
                    #     print(curr_obs[key].dtype)

                    if args.data_config == "oxe_droid" and action_space == "cartesian_position":
                        request_data = {
                            "video.exterior_image_1": resized_exterior_image,
                            "video.exterior_image_2": mask_3ch,
                            "video.wrist_image":  resized_wrist_image,
                            "state.eef_position": curr_obs["eef_position"],
                            "state.eef_rotation": curr_obs["eef_rotation"],
                            "state.gripper_position": curr_obs["gripper_position"],
                            "annotation.language.language_instruction": instruction,
                        }
                    elif args.data_config == "oxe_droid" and action_space == "joint_position":
                        request_data = {
                            "video.exterior_image_1": resized_exterior_image,
                            "video.exterior_image_2": mask_3ch,
                            "video.wrist_image":  resized_wrist_image,
                            "state.eef_position": curr_obs["eef_position"],
                            "state.eef_rotation": curr_obs["eef_rotation"],
                            "state.gripper_position": curr_obs["gripper_position"],
                            "annotation.language.language_instruction": instruction,
                        }
                    elif args.data_config == "real_droid_cartesian_position" and action_space == "cartesian_position":
                        request_data = {
                            "video.exterior_image_1_left": curr_obs['left_image'],
                            "video.wrist_image_left":  curr_obs['wrist_image'],
                            "state.end_effector_position": curr_obs["eef_position"],
                            "state.end_effector_rotation": curr_obs["eef_rotation"],
                            "state.gripper_close": curr_obs["gripper_position"],
                            "annotation.human.action.task_description": instruction,
                        }
                        
                    elif args.data_config == "real_droid_joint_position" and action_space == "joint_position":
                        # Add the leading time axis the modality transforms expect:
                        # video -> (T, H, W, C), state -> (T, D), language -> list of len T.
                        request_data = {
                            "video.exterior_image_1_left": curr_obs['left_image'][None],
                            "video.wrist_image_left":  curr_obs['wrist_image'][None],
                            "state.joint_pos_abs": curr_obs["joint_position"][None],
                            "state.gripper_close": curr_obs["gripper_position"][None],
                            "annotation.human.action.task_description": [instruction],
                        }

                    elif args.data_config == "real_droid_joint_3view" and action_space == "joint_position":
                        # Play-teleop 3-view checkpoints (periphanes/gr00t-n15-play-teleop-*-3view-flare-bsz32).
                        # Camera -> key mapping follows the training data:
                        #   exterior_image_1_left = ZED 34022131, exterior_image_2_left = ZED 37149196,
                        #   wrist_image_left = ZED 11845904 (all left lenses, native 720x1280).
                        request_data = {
                            "video.exterior_image_1_left": exterior1_image[None],
                            "video.exterior_image_2_left": curr_obs['left_image'][None],
                            "video.wrist_image_left":  curr_obs['wrist_image'][None],
                            "state.joint_pos_abs": curr_obs["joint_position"][None],
                            "state.gripper_close": curr_obs["gripper_position"][None],
                            "annotation.human.action.task_description": [instruction],
                        }
                        
                        
                    # pick up the blue cube and put into the white basket.
                    else: 
                        print("[Error] Not supported data_config or action_space")
                        sys.exit(1)
                        return
                    

                    # Wrap the server call in a context manager to prevent Ctrl+C from interrupting it
                    # Ctrl+C will be handled after the server call is complete
                    # obs_time = time.time() - start_time
    
                    with prevent_keyboard_interrupt():
                        infer_start = time.time()
                        # this returns action chunk [10, 8] of 10 joint velocity actions (7) + gripper position (1)
                        action = policy.get_action(request_data)
                        keys = list(action.keys())
                        print(action)
                        print(keys)
                        
                        if args.data_config == "real_droid_cartesian_position" and action_space == "cartesian_position":
                            concatenated_action = np.concatenate([
                                action[keys[0]], 
                                action[keys[1]], 
                                action[keys[2]].reshape(-1, 1)  # 1차원을 2차원으로 변환
                            ], axis=1) 
                            
                        elif args.data_config in ("real_droid_joint_position", "real_droid_joint_3view") and action_space == "joint_position":
                            concatenated_action = np.concatenate([
                                action[keys[0]], 
                                action[keys[1]].reshape(-1, 1)  # 1차원을 2차원으로 변환
                            ], axis=1)
                        
                             
                        print("output: ", output)
                        txt.append(output)
                        
                        infer_end = time.time()
                    
                    
                    inference_time = infer_end - infer_start
                    inference_times.append(inference_time)
                    
                    #assert pred_action_chunk.shape == (10, 8)

                    # === action chunk별 process 기록 ===
                    with open(process_tracking_tsv, "a", newline="") as f:
                        writer = csv.writer(f, delimiter="\t")
                        writer.writerow([
                            episode_id,
                            action_chunk_idx,
                            #json.dumps(score_dict, ensure_ascii=False),
                            #json.dumps(best_score_dict, ensure_ascii=False),
                            "",  # episode_score_sum
                            "",  # episode_action_chunk_sum
                            #json.dumps(pred_action_chunk.tolist(), ensure_ascii=False),
                            ""
                        ])
                    if args.save_trajectories:
                        traj["pred_chunk"].append(np.asarray(concatenated_action, dtype=np.float32))
                        traj["pred_chunk_t_step"].append(t_step)
                        traj["pred_chunk_infer_time"].append(inference_time)

                    action_chunk_idx += 1
                    #all_best_scores_dict.append(best_score_dict) 

                # Select current action to execute from chunk
                action = concatenated_action[actions_from_chunk_completed]
                actions_from_chunk_completed += 1

                # Binarize gripper action
                #print(action)
                # if action[-1].item() > 0.5:
                #     # action[-1] = 1.0
                #     #action = np.concatenate([action[:-1], np.ones((1,))])
                #     action = np.concatenate([action[:-1], ])
                # else:
                #     # action[-1] = 0.0
                #     action = np.concatenate([action[:-1], np.zeros((1,))])
                
                # clip all dimensions of action to [-1, 1]
                if "velocity" in action_space:
                    action = np.clip(action, -1, 1)
                
                # Safety clamp between the decoded chunk and the controller (see
                # _clamp_action). Applied to nominal and perturbed runs alike so the
                # clamp itself is never a difference between branches.
                action_preclamp = np.asarray(action, dtype=np.float32).copy()
                n_clamped, max_excess = 0, 0.0
                if args.clamp_actions:
                    action, n_clamped, max_excess = _clamp_action(
                        action, curr_obs["joint_position"],
                        args.max_joint_step, args.gripper_range,
                    )
                    if n_clamped:
                        clamp_events += 1
                        clamp_dims_total += n_clamped
                        clamp_max_excess = max(clamp_max_excess, max_excess)

                if args.save_trajectories:
                    traj["action_executed"].append(np.asarray(action, dtype=np.float32))
                    traj["action_preclamp"].append(action_preclamp)
                    traj["clamped_dims"].append(np.int16(n_clamped))
                    traj["chunk_idx"].append(action_chunk_idx - 1)

                env.step(action)
                # print(action)
                # Sleep to match DROID data collection frequency
                elapsed_time = time.time() - start_time
                if elapsed_time < 1 / DROID_CONTROL_FREQUENCY:
                    time.sleep(1 / DROID_CONTROL_FREQUENCY - elapsed_time)
            except KeyboardInterrupt:
                interrupted = True
                break
            except Exception as _exc:
                # Anything else -- most often the robot control server dropping its
                # zerorpc heartbeat mid-rollout ("LostRemote after 80s heartbeat"), which
                # kills get_observation(). Without this branch the exception unwinds past
                # _save_episode and every recorded step is lost: the per-camera mp4s are
                # left on disk as an orphaned prefix with no manifest and no states.
                import traceback
                crashed_error = f"{type(_exc).__name__}: {_exc}"
                print(f"\n[crash] rollout aborted at step {t_step}: {crashed_error}")
                traceback.print_exc()
                break

        if crashed_error is not None:
            # Save the prefix before doing anything else, then stop: a dead control server
            # will not recover on its own, and continuing would start an episode against
            # a robot in an unknown state.
            npz_path = _save_episode(
                args, episodes_dir, episode_id, traj, instruction,
                success=None, cancelled=True, interrupted=True,
                duration=time.time() - start_rollout_time,
                vid_writer=vid_writer,
                velocity_noise=velocity_noise,
                clamp_stats=_clamp_stats(),
                tuple_id=args.tuple_id,
                attempt_index=attempt_index,
                crashed=crashed_error,
            )
            print(f"[crash] partial trajectory saved: {npz_path}")
            print("[crash] the robot control server is most likely down -- restart it "
                  "(and check the arm is not in a reflex/error state) before rerunning.")
            raise SystemExit(1)

        if interrupted:
            cancel = input("Do you want to cancel this execution? (y/n): ").lower()
            if cancel in ["y"]:
                print("This execution will not be included in the statistics.")
                # Still persist the trajectory -- it is excluded from the success
                # statistics, not discarded. success=None, cancelled=True.
                _save_episode(
                    args, episodes_dir, episode_id, traj, instruction,
                    success=None, cancelled=True, interrupted=True,
                    duration=time.time() - start_rollout_time,
                    vid_writer=vid_writer,
                    velocity_noise=velocity_noise,
                    clamp_stats=_clamp_stats(),
                    tuple_id=args.tuple_id,
                    attempt_index=attempt_index,
                )
                if input("Do one more eval? (enter y or n) ").lower() not in ["y", ""]: # "" = enter
                    break
                last_instruction = instruction
                env.reset()
                while input("Does robot reset? (enter y or n) ").lower() in ["n"]:
                    env.reset()
                continue
        
        # save the outputs of the inference: 
        save_outputs = os.path.join(results_dir, f"outputs_{num_trials+1:03d}.json")
        with open(save_outputs, 'w') as f:
            json.dump(txt, f, indent=2, default=str) 

        if args.preview_video and len(video) > 0:
            video = np.stack(video)
            video_filename = f"{num_trials+1:03d}.mp4"
            save_filename = os.path.join(videos_dir, video_filename)
            # fps must match the control frequency; this used to be hardcoded to 15 while
            # the loop runs at 10 Hz, so the preview played 1.5x fast.
            # Frames were already downscaled at append time.
            ImageSequenceClip(list(video), fps=args.traj_video_fps).write_videofile(
                save_filename, codec="libx264"
            )

        duration = time.time() - start_rollout_time
        success = get_success_rate()
        _save_episode(
            args, episodes_dir, episode_id, traj, instruction,
            success=success, cancelled=False, interrupted=interrupted,
            duration=duration,
            vid_writer=vid_writer,
            velocity_noise=velocity_noise,
            clamp_stats=_clamp_stats(),
            tuple_id=args.tuple_id,
            attempt_index=attempt_index,
        )
        success_list.append(success)
        num_trials += 1
        if success == 0.5:
            grasp += 1
        elif success == 1.0:
            success_cnt += 1
            grasp += 1
        grasp_rate = grasp / len(success_list)
        success_rate = success_cnt / len(success_list)
        avg_inference_time = np.mean(inference_times) if inference_times else 0
        
        # === 에피소드별 score sum, action chunk sum 기록 ===
        episode_score_sum_dict = {}
        # action chunk별 best_score_dict를 누적합
        # (action chunk별 기록에서 all_best_scores_dict에 append 필요)
        for d in all_best_scores_dict:
            for k, v in d.items():
                episode_score_sum_dict[k] = episode_score_sum_dict.get(k, 0) + (v if v is not None else 0)
        with open(process_tracking_tsv, "a", newline="") as f:
            writer = csv.writer(f, delimiter="\t")
            writer.writerow([
                episode_id,
                "episode_summary",
                "",
                "",
                json.dumps(episode_score_sum_dict, ensure_ascii=False),
                action_chunk_idx,
                "",
                success  # 에피소드 성공 여부 추가
            ])
        # task별로 inference time 누적
        task_key = instruction.lower().strip()
        task_inference_times[task_key].append(avg_inference_time)
        print(f"Current trajectory success: {success}")
        print(f"Overall success rate: {success_rate:.3f}")
        print(f"Overall grasp rate: {grasp_rate:.3f}")
        print(f"Average inference time (this trial): {avg_inference_time:.4f} sec")
        df = df[df["task"].str.lower().str.strip() != task_key]
        df.loc[len(df)] = [
            instruction,
            grasp,
            success_cnt,
            f"{grasp_rate * 100:.3f}", 
            f"{success_rate * 100:.3f}",    
            num_trials,
            f"{avg_inference_time:.4f}",
        ]
        df.to_csv(tsv_filename, sep='\t', index=False)
        if input("Do one more eval? (enter y or n) ").lower() not in ["y", ""]:
            break
        last_instruction = instruction
        env.reset()
        while input("Does robot reset? (enter y or n) ").lower() in ["n"]:
            env.reset()
        
    # 각 task별로 avg_inference_time 갱신
    for idx, row in df.iterrows():
        task_key = row["task"].lower().strip()
        avg_time = np.mean(task_inference_times[task_key]) if task_inference_times[task_key] else 0
        df.at[idx, "avg_inference_time"] = f"{avg_time:.4f}"
    df.to_csv(tsv_filename, sep='\t', index=False)
    print(f"Results saved to {tsv_filename}")


def _resize_for_traj(img, size: int):
    """Downscale a camera frame for trajectory storage. size <= 0 keeps native res."""
    if size and size > 0:
        import cv2

        img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA)
    return np.ascontiguousarray(img, dtype=np.uint8)


def _clamp_action(action, joint_pos, max_joint_step: float, gripper_range):
    """Clamp a commanded action against the MEASURED joint position.
    Returns (clamped_action, num_clamped_dims, max_excess_rad).
    The action space here is 7 absolute joint targets (rad) + gripper. Each joint target
    is limited to `max_joint_step` radians from where the arm currently *is*, which is
    the property that actually bounds the motion: limiting the step relative to the
    previous *command* would let a chunk walk away from a lagging arm one safe-looking
    step at a time.
    Every clamp is counted and reported, because a silently clamped perturbation is a
    perturbation that did not happen -- if the clamp rate is high, the recorded sigma is
    a fiction and the episode's divergence is really measuring the clamp.
    """
    action = np.asarray(action, dtype=np.float32).copy()
    n_joints = min(7, action.shape[0], np.asarray(joint_pos).shape[0])
    ref = np.asarray(joint_pos, dtype=np.float32)[:n_joints]
    target = action[:n_joints]
    lo, hi = ref - max_joint_step, ref + max_joint_step
    excess = np.maximum(np.abs(target - ref) - max_joint_step, 0.0)
    num_clamped = int((excess > 0).sum())
    max_excess = float(excess.max()) if excess.size else 0.0
    action[:n_joints] = np.clip(target, lo, hi)
    if action.shape[0] > 7:
        action[7] = float(np.clip(action[7], gripper_range[0], gripper_range[1]))
    return action, num_clamped, max_excess


def _downscale_preview(img, scale: float):
    """Scale a preview frame by `scale`, rounding to even dimensions for libx264."""
    if not scale or scale >= 1.0:
        return np.ascontiguousarray(img, dtype=np.uint8)
    import cv2

    h, w = img.shape[:2]
    nw = max(2, int(round(w * scale)) // 2 * 2)
    nh = max(2, int(round(h * scale)) // 2 * 2)
    return np.ascontiguousarray(
        cv2.resize(img, (nw, nh), interpolation=cv2.INTER_AREA), dtype=np.uint8
    )


def _parse_wh(spec: str) -> tuple[int, int]:
    """Parse a "WxH" resolution spec -> (width, height)."""
    w, _, h = str(spec).lower().partition("x")
    return int(w), int(h)


def _crop_resize_for_video(img, width: int, height: int):
    """Center-crop `img` to the target aspect ratio, then resize to (width, height).
    Cropping before resizing keeps the scene undistorted. Squashing 1280x720 (16:9)
    straight into 256x192 (4:3) would stretch it vertically by 1.33x, and a world model
    trained on stretched frames cannot be evaluated against anything else.
    """
    import cv2

    h, w = img.shape[:2]
    target_ar = width / height
    if w / h > target_ar:  # too wide -> trim the sides
        new_w = int(round(h * target_ar))
        x0 = (w - new_w) // 2
        img = img[:, x0 : x0 + new_w]
    elif w / h < target_ar:  # too tall -> trim top/bottom
        new_h = int(round(w / target_ar))
        y0 = (h - new_h) // 2
        img = img[y0 : y0 + new_h, :]
    img = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
    return np.ascontiguousarray(img, dtype=np.uint8)


class _PerCameraVideoWriter:
    """Incremental h264 writers, one mp4 per camera, for one episode.
    Frames are encoded as they arrive, so RAM stays flat regardless of episode length
    and a crash or Ctrl+C still leaves a playable prefix on disk. Keys are the GR00T
    model video keys (video.exterior_image_1_left, ...) so the produced tree is already
    in DreamDojo's `videos/<video_key>/` shape.
    """

    def __init__(self, out_dir: str, stem: str, size: str, fps: float):
        import imageio.v2 as imageio

        self._imageio = imageio
        self.width, self.height = _parse_wh(size)
        self.fps = float(fps)
        self.out_dir = out_dir
        self.stem = stem
        self._writers: dict[str, object] = {}
        self._counts: dict[str, int] = {}
        self.paths: dict[str, str] = {}

    def _writer_for(self, video_key: str):
        if video_key not in self._writers:
            cam_dir = os.path.join(self.out_dir, video_key)
            os.makedirs(cam_dir, exist_ok=True)
            path = os.path.join(cam_dir, self.stem + ".mp4")
            # macro_block_size=None: 256x192 are multiples of 16 already, and letting
            # ffmpeg pad to its default macro block would silently change the resolution.
            self._writers[video_key] = self._imageio.get_writer(
                path, fps=self.fps, codec="libx264", quality=8,
                macro_block_size=None, ffmpeg_log_level="error",
            )
            self._counts[video_key] = 0
            self.paths[video_key] = path
        return self._writers[video_key]

    def append(self, video_key: str, img):
        self._writer_for(video_key).append_data(
            _crop_resize_for_video(img, self.width, self.height)
        )
        self._counts[video_key] += 1

    def close(self) -> dict:
        """Close every writer and return {video_key: {path, num_frames}}."""
        info = {}
        for key, w in self._writers.items():
            with contextlib.suppress(Exception):
                w.close()
            info[key] = {
                "path": os.path.relpath(self.paths[key], self.out_dir),
                "num_frames": self._counts[key],
            }
        self._writers.clear()
        return info


def _save_episode(
    args: Args,
    episodes_dir: str,
    episode_id: int,
    traj: dict,
    instruction: str,
    success,
    cancelled: bool,
    interrupted: bool,
    duration: float,
    vid_writer: "_PerCameraVideoWriter | None" = None,
    velocity_noise=None,
    clamp_stats: dict | None = None,
    tuple_id: str | None = None,
    attempt_index: int | None = None,
    crashed: str | None = None,
):
    """Write one episode's trajectory to <episodes_dir> as an .npz + .json pair.
    Called for EVERY episode -- successes, partials, failures and Ctrl+C cancels --
    so failure trajectories are retained. `success` is None for cancelled episodes.
    """
    if not args.save_trajectories:
        return None

    # Close the per-camera mp4s FIRST, unconditionally: this runs on the cancel path and
    # the normal path alike, and an unclosed imageio writer leaves a truncated file.
    video_info = vid_writer.close() if vid_writer is not None else {}

    # Per-step buffers can differ in length by one if Ctrl+C landed mid-step
    # (observation recorded, action not yet). Truncate to the common prefix so
    # every per-step array stays index-aligned.
    per_step_keys = [
        "exterior_image",
        "exterior1_image",
        "wrist_image",
        "joint_pos_abs",
        "gripper_close",
        "eef_position",
        "eef_rotation",
        "t_step",
        "time_rel",
        "action_executed",
        "action_preclamp",
        "clamped_dims",
        "chunk_idx",
    ]
    present = [k for k in per_step_keys if len(traj.get(k, [])) > 0]
    if not present:
        print("[traj] nothing recorded for this episode, skipping save")
        return None
    n_steps = min(len(traj[k]) for k in present)

    arrays = {k: np.stack(traj[k][:n_steps]) for k in present}
    for k in ("pred_chunk", "pred_chunk_t_step", "pred_chunk_infer_time"):
        if len(traj.get(k, [])) > 0:
            arrays[k] = np.stack(traj[k]) if k == "pred_chunk" else np.asarray(traj[k])

    if cancelled:
        label = "cancelled"
    else:
        label = f"success{success}"
    # attempt_index, not episode_id: see the note at its increment. Falls back to
    # episode_id+1 only if a caller predates the counter.
    stem = f"ep_{attempt_index if attempt_index is not None else episode_id + 1:03d}_{label}"
    npz_path = os.path.join(episodes_dir, stem + ".npz")
    json_path = os.path.join(episodes_dir, stem + ".json")

    saver = np.savez_compressed if args.traj_compress else np.savez
    saver(npz_path, **arrays)

    meta = {
        "episode_id": episode_id,
        "attempt_index": attempt_index,
        # Set when the episode ended in an unhandled exception (e.g. the robot server
        # dropping its zerorpc heartbeat) rather than completing or being cancelled.
        # The trajectory is a valid prefix, but it is NOT a completed episode.
        "crashed": crashed,
        "instruction": instruction,
        "success": success,
        "cancelled": cancelled,
        "interrupted": interrupted,
        "num_steps": int(n_steps),
        "duration_sec": round(float(duration), 3),
        "arrays": {k: list(v.shape) for k, v in arrays.items()},
        # --- world-model (action_sensitive_wm) contract ---
        # Everything below is what the LeRobot/DreamDojo converter needs so it never has
        # to guess a layout. See convert_real_droid_to_lerobot.py in the WM repo.
        "wm": {
            "fps": float(args.traj_video_fps),
            # observation.state and action are built by concatenating these keys IN ORDER.
            # The order mirrors RealDroidJointDataConfig.state_keys / .action_keys, i.e.
            # exactly what the policy was trained on -- do not reorder.
            "state_keys": [
                {"name": "state.joint_pos_abs", "array": "joint_pos_abs", "dim": 7},
                {"name": "state.gripper_close", "array": "gripper_close", "dim": 1},
            ],
            "action_keys": [
                {"name": "action.joint_pos_abs", "slice": [0, 7], "dim": 7},
                {"name": "action.gripper_close", "slice": [7, 8], "dim": 1},
            ],
            "action_array": "action_executed",
            "action_dim": 8,
            "state_dim": 8,
            "action_space": "joint_position_absolute",
            "videos": video_info,
            "video_size": [vid_writer.width, vid_writer.height] if vid_writer else None,
            "video_fps": float(args.traj_video_fps),
            "video_note": (
                "center-cropped to the target aspect ratio, then resized; NOT squashed. "
                "NOT the exact policy input, which is VideoCropDroid-cropped then 224x224."
            ),
            # Counterfactual-tuple fields (WM avenue 3 / provenance.jsonl).
            # With no perturbation these stay "control"/0.0 and each episode is its own
            # singleton tuple. With --noise-sigma>0 and a fixed --noise-init-seed, the
            # episodes sharing a --tuple-id are branches of one same-start tuple: same
            # initial sampler noise, different perturbation stream.
            "variant": velocity_noise.variant_name() if velocity_noise else "control",
            "sigma": float(velocity_noise.sigma) if velocity_noise else 0.0,
            "perturbed": bool(velocity_noise and velocity_noise.sigma),
            "noise_options": velocity_noise.effective_options() if velocity_noise else None,
            "tuple_id": tuple_id,
            "k_index": episode_id,
            # Clamp telemetry. A high clamp rate means the recorded sigma is a fiction:
            # the divergence being measured is the clamp's, not the perturbation's.
            "clamp": clamp_stats,
            # For post-hoc start-state matching (FR3 plan sec.4-C): the measured start state,
            # which is the only start_key a real robot can offer.
            "start_joint_pos": (
                arrays["joint_pos_abs"][0].tolist() if "joint_pos_abs" in arrays else None
            ),
            "start_gripper": (
                arrays["gripper_close"][0].tolist() if "gripper_close" in arrays else None
            ),
            "start_eef_position": (
                arrays["eef_position"][0].tolist() if "eef_position" in arrays else None
            ),
            "start_eef_rotation": (
                arrays["eef_rotation"][0].tolist() if "eef_rotation" in arrays else None
            ),
        },
        "traj_video_mode": args.traj_video_mode,
        "traj_image_size": args.traj_image_size,
        "traj_image_note": (
            "camera frames resized to a square; NOT the exact model input, which is "
            "VideoCropDroid-cropped then resized to 224x224"
            if args.traj_image_size and args.traj_image_size > 0
            else "native camera resolution"
        ),
        "checkpoint_dir": getattr(args.policy, "dir", None),
        "policy_config": getattr(args.policy, "config", None),
        "data_config": args.data_config,
        "embodiment_tag": args.embodiment_tag,
        "denoising_steps": args.denoising_steps,
        "vision_token_num": args.vision_token_num,
        "open_loop_horizon": args.open_loop_horizon,
        "max_timesteps": args.max_timesteps,
        "control_frequency_hz": DROID_CONTROL_FREQUENCY,
        "left_camera_id": args.left_camera_id,
        "wrist_camera_id": args.wrist_camera_id,
        "exterior1_camera_id": args.exterior1_camera_id if "exterior1_image" in arrays else None,
        "camera_to_model_key": (
            {
                "video.exterior_image_1_left": args.exterior1_camera_id,
                "video.exterior_image_2_left": args.left_camera_id,
                "video.wrist_image_left": args.wrist_camera_id,
            }
            if args.data_config == "real_droid_joint_3view"
            else {
                "video.exterior_image_1_left": args.left_camera_id,
                "video.wrist_image_left": args.wrist_camera_id,
            }
        ),
        "saved_at": datetime.datetime.now().isoformat(timespec="seconds"),
    }
    with open(json_path, "w") as f:
        json.dump(meta, f, indent=2, default=str)

    size_mb = os.path.getsize(npz_path) / 1e6
    print(f"[traj] saved {stem} ({n_steps} steps, {size_mb:.1f} MB) -> {episodes_dir}")
    return npz_path


def _extract_observation(args: Args, obs_dict):
    image_observations = obs_dict["image"]
    left_image, wrist_image, exterior1_image = None, None, None

    for key in image_observations:
        if "left" not in key:
            continue
        
        # Note the "left" below refers to the left camera in the stereo pair.
        # The model is only trained on left stereo cams, so we only feed those.
        if args.left_camera_id in key:
            left_image = image_observations[key]
        elif args.wrist_camera_id in key:
            wrist_image = image_observations[key]
        elif args.exterior1_camera_id in key:
            # Second exterior ZED; only consumed by 3-view data configs.
            exterior1_image = image_observations[key]
    
    if left_image is None or wrist_image is None:
        raise RuntimeError(
            f"Missing camera frames: left={args.left_camera_id} found={left_image is not None}, "
            f"wrist={args.wrist_camera_id} found={wrist_image is not None}; "
            f"available keys={list(image_observations.keys())}"
        )

    # Drop the alpha dimension, convert BGR -> RGB
    left_image = left_image[..., :3][..., ::-1]
    wrist_image = wrist_image[..., :3][..., ::-1]
    if exterior1_image is not None:
        exterior1_image = exterior1_image[..., :3][..., ::-1].astype(np.uint8)

    # In addition to image observations, also capture the proprioceptive state
    robot_state = obs_dict["robot_state"]
    joint_position = np.array(robot_state["joint_positions"])
    gripper_position = np.array([robot_state["gripper_position"]])
    cartesian_position = np.array(robot_state["cartesian_position"][:3])
    cartesian_rotation = np.array(robot_state["cartesian_position"][3:])
    

    return {
        "left_image": left_image.astype(np.uint8),
        "wrist_image": wrist_image.astype(np.uint8) ,
        "exterior1_image": exterior1_image,
        "joint_position": joint_position,
        "eef_position": cartesian_position,
        "eef_rotation": cartesian_rotation, 
        "gripper_position": gripper_position,
    }
    



if __name__ == "__main__":
    args: Args = tyro.cli(Args)
    main(args)



    
# Tasks 
# 1. Pick up the top red cup from the stack and place it next to the blue cup.
# 2. Pick up the yellow sponge gently and place it in the gold bowl.
# 3. Use the red eraser to clean the board.
# 4. Pick up the yellow charger and plug it into the socket.