import argparse import time import os # from wall_x.utils.baseline_utils import check_baseline_dump, update_baseline from wall_x.infer.utils_libero import set_seed_everywhere, TaskSuite, TASK_MAX_STEPS from wall_x.infer.infer_config import InferConfig from wall_x.infer.env_libero import LiberoRobotEnv if __name__ == "__main__": args = argparse.ArgumentParser(description="Wall-X Libero evaluation script") args.add_argument("--seed", type=int, default=42, help="Random seed") args.add_argument("--id", type=int, default=None, help="Unique index id") args.add_argument("--name", type=int, default=None, help="Launch command name") args.add_argument( "--baseline_path", type=str, default=None, help="Path to baseline record table" ) args.add_argument( "--update_baseline", type=bool, default=False, help="Whether to update baseline table", ) args.add_argument( "--mode", type=str, default="flow", choices=["flow", "ar"], help="Running mode" ) args.add_argument( "--checkpoint_path", type=str, required=True, help="Model checkpoint path" ) args.add_argument( "--train_config_path", type=str, required=False, default=None, help="Path to training config .yml file", ) args.add_argument( "--norm_key", type=str, default="physical-intelligence/libero", help="Key for normalization statistics", ) args.add_argument( "--cam_names", nargs="+", default=["face_view", "right_wrist_view"], help="List of camera names (e.g., --cam_names face_view right_wrist_view)", ) args.add_argument( "--task_suite_name", type=str, default=TaskSuite.LIBERO_SPATIAL, choices=[e.value for e in TaskSuite], help="Libero task suite to load", ) args.add_argument( "--initial_states_path", type=str, default="DEFAULT", help="Path to initial states .json file, or 'DEFAULT' to use default states.", ) args.add_argument( "--num_trials_per_task", type=int, default=50, help="Number of evaluation episodes to run per task", ) args.add_argument( "--rollout_dir", type=str, default="./rollouts", help="Directory to save rollout videos", ) args = args.parse_args() print(f"Using random seed: {args.seed}") set_seed_everywhere(args.seed) print("Initializing InferConfig...") if args.train_config_path is None: args.train_config_path = os.path.join(args.checkpoint_path, "config.yml") config = InferConfig( checkpoint_path=args.checkpoint_path, train_config_path=args.train_config_path, norm_key=args.norm_key, cam_names=args.cam_names, ) if args.mode == "flow": config.action_horizon = config.train_config.get("data", {}).get( "action_horizon_flow", 10 ) elif args.mode == "ar": config.action_horizon = config.train_config.get("data", {}).get( "action_horizon_ar", 10 ) else: raise ValueError(f"Invalid mode: {args.mode}") config.model_device = "cuda" print("Initializing LiberoRobotEnv (Evaluator)...") config.action_dim = 7 config.pred_horizon = 10 evaluator = LiberoRobotEnv( config=config, task_suite_name=args.task_suite_name, initial_states_path=args.initial_states_path, rollout_dir=args.rollout_dir, seed=args.seed, ) print(f"\n{'='*20} Starting Evaluation {'='*20}") print(f"Task suite: {args.task_suite_name}") print(f"Number of tasks: {evaluator.num_tasks}") print(f"Trials per task: {args.num_trials_per_task}") print(f"Initial states: {args.initial_states_path}") print(f"Videos will be saved to: {evaluator.rollout_dir}") print(f"{'='*50}\n") total_successes = 0 total_episodes_run = 0 start_time = time.time() for task_id in range(evaluator.num_tasks): task_successes = 0 task_episodes_attempted = 0 libero_env_instance = None task_desc = "" initial_states = None max_infer_times = TASK_MAX_STEPS[args.task_suite_name] print( f"{args.task_suite_name} TASK_MAX_STEPS: {TASK_MAX_STEPS[args.task_suite_name]}" ) for ep_idx in range(args.num_trials_per_task): print(f" > Running trial {ep_idx + 1} / {args.num_trials_per_task}...") try: print(f"\nCreating environment for Task {task_id}...") libero_env_instance, task_desc, initial_states = ( evaluator.create_env_for_task(task_id) ) print( f"--- Starting task {task_id + 1} / {evaluator.num_tasks}: {task_desc} ---" ) except Exception as e: print( f"\n[CRITICAL ERROR] Failed to create environment for task {task_id}: {e}. Skipping entire task." ) continue task_episodes_attempted += 1 total_episodes_run += 1 success = False try: if args.mode == "flow": success = evaluator.run_infer_flow_action( env=libero_env_instance, task_id=task_id, task_desc=task_desc, default_initial_states=initial_states, episode_idx=ep_idx, max_infer_times=max_infer_times, ) elif args.mode == "ar": success = evaluator.run_infer_ar_action( env=libero_env_instance, task_id=task_id, task_desc=task_desc, default_initial_states=initial_states, episode_idx=ep_idx, max_infer_times=max_infer_times, ) except Exception as e: print(f" [EXCEPTION] Episode run error: {e}") if success: task_successes += 1 total_successes += 1 print(" > Trial result: SUCCESS") else: print(" > Trial result: FAILURE") if task_episodes_attempted > 0: print( f" > Task {task_id} current success rate: {task_successes / task_episodes_attempted * 100:.1f}% ({task_successes}/{task_episodes_attempted})" ) if total_episodes_run > 0: print( f" > Overall current success rate: {total_successes / total_episodes_run * 100:.1f}% ({total_successes}/{total_episodes_run})" ) task_success_rate = ( task_successes / task_episodes_attempted if task_episodes_attempted > 0 else 0 ) print(f"\n--- Task {task_id} ({task_desc}) Summary ---") print( f"Success rate: {task_success_rate * 100:.1f}% ({task_successes}/{task_episodes_attempted})" ) print(f"{'-'*40}\n") end_time = time.time() total_time = end_time - start_time final_success_rate = ( total_successes / total_episodes_run if total_episodes_run > 0 else 0 ) print(f"\n{'='*20} Final Evaluation Summary {'='*20}") print(f"Total runtime: {total_time:.2f} seconds ({total_time / 60:.1f} minutes)") print(f"Total trials run: {total_episodes_run}") print(f"Total successes: {total_successes}") print(f"Overall success rate: {final_success_rate * 100:.2f}%") print(f"{'='*56}") print("Evaluation completed.")