| |
| |
| |
| |
| @@ -80,6 +80,9 @@ class WandBConfig: |
| @dataclass |
| class EvalConfig: |
| n_episodes: int = 50 |
| + # Number of rollout videos to render. Set to 0 to disable lightweight MP4 |
| + # rendering. Full dataset recording (`recording=true`) takes precedence. |
| + max_episodes_rendered: int = 10 |
| # `batch_size` specifies the number of environments to use in a gym.vector.VectorEnv. |
| # Set to 0 for auto-tuning based on available CPU cores and n_episodes. |
| batch_size: int = 0 |
| @@ -95,6 +98,9 @@ class EvalConfig: |
| recording_private: bool = False |
| |
| def __post_init__(self) -> None: |
| + if self.max_episodes_rendered < 0: |
| + raise ValueError("eval.max_episodes_rendered must be non-negative.") |
| + self.max_episodes_rendered = min(self.max_episodes_rendered, self.n_episodes) |
| if self.recording_repo_id is not None and not self.recording: |
| raise ValueError("eval.recording_repo_id requires eval.recording=true.") |
| if self.batch_size == 0: |
| |
| |
| |
| |
| @@ -38,6 +38,22 @@ ACTION_DIM = 12 # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gr |
| ACTION_LOW = -1.0 |
| ACTION_HIGH = 1.0 |
| |
| +STAGE_REACHED_HANDLE = "reached_handle" |
| +STAGE_VALID_GRASP = "valid_grasp" |
| +STAGE_DOOR_STARTED_MOVING = "door_started_moving" |
| +STAGE_GRASP_MAINTAINED = "grasp_maintained" |
| +STAGE_METRIC_KEYS = ( |
| + STAGE_REACHED_HANDLE, |
| + STAGE_VALID_GRASP, |
| + STAGE_DOOR_STARTED_MOVING, |
| + STAGE_GRASP_MAINTAINED, |
| +) |
| + |
| +HANDLE_REACH_DISTANCE_M = 0.08 |
| +DOOR_MOVEMENT_THRESHOLD = 0.05 |
| +DOOR_TARGET_THRESHOLD = 0.90 |
| +GRASP_CONTACT_GRACE_STEPS = 5 |
| + |
| # Default PandaOmron cameras. We surface these raw names directly as |
| # `observation.images.<name>` so the LeRobot dataset/policy keys match |
| # RoboCasa's native convention (no implicit renaming). |
| @@ -160,6 +176,7 @@ class RoboCasaEnv(gym.Env): |
| # to avoid inheriting stale GPU/EGL contexts across fork(). |
| self._env: Any = None |
| self.task_description = "" |
| + self._stage_state: dict[str, Any] = {} |
| |
| images = { |
| cam: spaces.Box( |
| @@ -248,6 +265,120 @@ class RoboCasaEnv(gym.Env): |
| assert self._env is not None |
| return self._env.render() |
| |
| + def _target_handle_geoms(self) -> list[str]: |
| + """Return collision geoms belonging to the current target cabinet handles.""" |
| + assert self._env is not None |
| + base_env = self._env.env |
| + fixture = getattr(base_env, "fxtr", None) |
| + if fixture is None: |
| + return [] |
| + |
| + prefix = fixture.name |
| + names: list[str] = [] |
| + for geom_id in range(base_env.sim.model.ngeom): |
| + name = base_env.sim.model.geom_id2name(geom_id) or "" |
| + if not name.startswith(prefix) or "door_handle" not in name.lower(): |
| + continue |
| + if "visual" in name.lower() or "reg_main" in name.lower(): |
| + continue |
| + names.append(name) |
| + return names |
| + |
| + def _reset_stage_state(self) -> None: |
| + """Initialize persistent per-episode OpenCabinet diagnostic state.""" |
| + assert self._env is not None |
| + base_env = self._env.env |
| + fixture = getattr(base_env, "fxtr", None) |
| + joint_names = list(getattr(fixture, "door_joint_names", [])) |
| + joint_state = fixture.get_joint_state(base_env, joint_names) if joint_names else {} |
| + self._stage_state = { |
| + **{key: False for key in STAGE_METRIC_KEYS}, |
| + "step": 0, |
| + "first_reach_step": None, |
| + "first_grasp_step": None, |
| + "first_motion_step": None, |
| + "last_grasp_step": None, |
| + "target_joints_with_grasp": set(), |
| + "initial_joint_state": joint_state, |
| + "max_door_openness": min(joint_state.values()) if joint_state else 0.0, |
| + "handle_geoms": self._target_handle_geoms(), |
| + } |
| + |
| + def _measure_stage_metrics(self) -> dict[str, Any]: |
| + """Measure and update the five-level OpenCabinet progress indicators.""" |
| + assert self._env is not None |
| + state = self._stage_state |
| + if self.task != "OpenCabinet" or not state: |
| + return {key: False for key in STAGE_METRIC_KEYS} |
| + |
| + base_env = self._env.env |
| + fixture = base_env.fxtr |
| + robot = base_env.robots[0] |
| + handle_geoms = state["handle_geoms"] |
| + step = int(state["step"]) |
| + |
| + eef_position = base_env.sim.data.site_xpos[robot.eef_site_id["right"]] |
| + handle_positions = [ |
| + base_env.sim.data.geom_xpos[base_env.sim.model.geom_name2id(name)] for name in handle_geoms |
| + ] |
| + handle_distance = ( |
| + min(float(np.linalg.norm(eef_position - position)) for position in handle_positions) |
| + if handle_positions |
| + else float("inf") |
| + ) |
| + reached_now = handle_distance <= HANDLE_REACH_DISTANCE_M |
| + |
| + gripper = robot.gripper.get("right") |
| + grasped_now = bool( |
| + gripper is not None and handle_geoms and base_env._check_grasp(gripper, handle_geoms) |
| + ) |
| + |
| + joint_names = fixture.door_joint_names |
| + joint_state = fixture.get_joint_state(base_env, joint_names) |
| + initial_state = state["initial_joint_state"] |
| + moved_now = any( |
| + joint_state[name] - initial_state.get(name, 0.0) >= DOOR_MOVEMENT_THRESHOLD |
| + for name in joint_names |
| + ) |
| + |
| + if reached_now and not state[STAGE_REACHED_HANDLE]: |
| + state["first_reach_step"] = step |
| + if grasped_now and not state[STAGE_VALID_GRASP]: |
| + state["first_grasp_step"] = step |
| + if moved_now and not state[STAGE_DOOR_STARTED_MOVING]: |
| + state["first_motion_step"] = step |
| + |
| + state[STAGE_REACHED_HANDLE] |= reached_now |
| + state[STAGE_VALID_GRASP] |= grasped_now |
| + state[STAGE_DOOR_STARTED_MOVING] |= moved_now |
| + if grasped_now: |
| + state["last_grasp_step"] = step |
| + |
| + recent_grasp = state["last_grasp_step"] is not None and ( |
| + step - state["last_grasp_step"] <= GRASP_CONTACT_GRACE_STEPS |
| + ) |
| + if recent_grasp: |
| + state["target_joints_with_grasp"].update( |
| + name for name, openness in joint_state.items() if openness >= DOOR_TARGET_THRESHOLD |
| + ) |
| + state[STAGE_GRASP_MAINTAINED] = bool(joint_names) and all( |
| + name in state["target_joints_with_grasp"] for name in joint_names |
| + ) |
| + state["max_door_openness"] = max( |
| + float(state["max_door_openness"]), |
| + min(joint_state.values()) if joint_state else 0.0, |
| + ) |
| + state["step"] = step + 1 |
| + |
| + return { |
| + **{key: bool(state[key]) for key in STAGE_METRIC_KEYS}, |
| + "handle_distance": handle_distance, |
| + "max_door_openness": float(state["max_door_openness"]), |
| + "first_reach_step": state["first_reach_step"] if state["first_reach_step"] is not None else -1, |
| + "first_grasp_step": state["first_grasp_step"] if state["first_grasp_step"] is not None else -1, |
| + "first_motion_step": state["first_motion_step"] if state["first_motion_step"] is not None else -1, |
| + } |
| + |
| def reset(self, seed=None, **kwargs): |
| self._ensure_env() |
| assert self._env is not None |
| @@ -262,6 +393,7 @@ class RoboCasaEnv(gym.Env): |
| |
| ep_meta = self._env.env.get_ep_meta() |
| self.task_description = ep_meta.get("lang", self.task) |
| + self._reset_stage_state() |
| |
| observation = self._format_raw_obs(raw_obs) |
| info = {"is_success": False} |
| @@ -279,9 +411,11 @@ class RoboCasaEnv(gym.Env): |
| action_dict = convert_action(action) |
| raw_obs, reward, done, truncated, info = self._env.step(action_dict) |
| |
| + stage_metrics = self._measure_stage_metrics() |
| + |
| is_success = bool(info.get("success", False)) |
| terminated = done or is_success |
| - info.update({"task": self.task, "done": done, "is_success": is_success}) |
| + info.update({"task": self.task, "done": done, "is_success": is_success, **stage_metrics}) |
| |
| observation = self._format_raw_obs(raw_obs) |
| if terminated: |
| @@ -289,6 +423,7 @@ class RoboCasaEnv(gym.Env): |
| "task": self.task, |
| "done": bool(done), |
| "is_success": bool(is_success), |
| + **stage_metrics, |
| } |
| self.reset() |
| |
| |
| |
| |
| |
| @@ -95,6 +95,13 @@ from lerobot.utils.utils import ( |
| inside_slurm, |
| ) |
| |
| +STAGE_METRIC_KEYS = ( |
| + "reached_handle", |
| + "valid_grasp", |
| + "door_started_moving", |
| + "grasp_maintained", |
| +) |
| + |
| |
| def _env_features_to_dataset_features(env_features: dict) -> dict: |
| """Convert EnvConfig.features to the dict format expected by LeRobotDataset.create().""" |
| @@ -245,6 +252,7 @@ def rollout( |
| all_rewards = [] |
| all_successes = [] |
| all_dones = [] |
| + all_stage_metrics: dict[str, list[Tensor]] = {key: [] for key in STAGE_METRIC_KEYS} |
| |
| step = 0 |
| # Keep track of which environments are done. |
| @@ -327,6 +335,13 @@ def rollout( |
| else: |
| successes = [False] * env.num_envs |
| |
| + for key in STAGE_METRIC_KEYS: |
| + value = info.get(key, np.zeros(env.num_envs, dtype=bool)) |
| + value_array = np.asarray(value, dtype=bool) |
| + if value_array.ndim == 0: |
| + value_array = np.full(env.num_envs, bool(value_array), dtype=bool) |
| + all_stage_metrics[key].append(torch.from_numpy(value_array)) |
| + |
| if recording_datasets is not None and raw_observation is not None: |
| prev_done = done.copy() |
| for env_idx in range(env.num_envs): |
| @@ -387,6 +402,7 @@ def rollout( |
| "reward": torch.stack(all_rewards, dim=1), |
| "success": torch.stack(all_successes, dim=1), |
| "done": torch.stack(all_dones, dim=1), |
| + **{key: torch.stack(values, dim=1) for key, values in all_stage_metrics.items()}, |
| } |
| if return_observations: |
| stacked_observations = {} |
| @@ -466,6 +482,7 @@ def eval_policy( |
| sum_rewards = [] |
| max_rewards = [] |
| all_successes = [] |
| + all_stage_metrics: dict[str, list[bool]] = {key: [] for key in STAGE_METRIC_KEYS} |
| all_seeds = [] |
| threads = [] # for video saving threads |
| n_episodes_rendered = 0 # for saving the correct number of videos |
| @@ -555,6 +572,9 @@ def eval_policy( |
| max_rewards.extend(batch_max_rewards.tolist()) |
| batch_successes = einops.reduce((rollout_data["success"] * mask), "b n -> b", "any") |
| all_successes.extend(batch_successes.tolist()) |
| + for key in STAGE_METRIC_KEYS: |
| + batch_stage = einops.reduce((rollout_data[key] * mask), "b n -> b", "any") |
| + all_stage_metrics[key].extend(batch_stage.tolist()) |
| if seeds: |
| all_seeds.extend(seeds) |
| else: |
| @@ -640,31 +660,39 @@ def eval_policy( |
| thread.join() |
| |
| # Compile eval info. |
| - info = { |
| - "per_episode": [ |
| + per_episode = [] |
| + for i, (sum_reward, max_reward, success, seed) in enumerate( |
| + zip( |
| + sum_rewards[:n_episodes], |
| + max_rewards[:n_episodes], |
| + all_successes[:n_episodes], |
| + all_seeds[:n_episodes], |
| + strict=True, |
| + ) |
| + ): |
| + per_episode.append( |
| { |
| "episode_ix": i, |
| "sum_reward": sum_reward, |
| "max_reward": max_reward, |
| "success": success, |
| "seed": seed, |
| + **{key: all_stage_metrics[key][i] for key in STAGE_METRIC_KEYS}, |
| } |
| - for i, (sum_reward, max_reward, success, seed) in enumerate( |
| - zip( |
| - sum_rewards[:n_episodes], |
| - max_rewards[:n_episodes], |
| - all_successes[:n_episodes], |
| - all_seeds[:n_episodes], |
| - strict=True, |
| - ) |
| - ) |
| - ], |
| + ) |
| + |
| + info = { |
| + "per_episode": per_episode, |
| "aggregated": { |
| "avg_sum_reward": float(np.nanmean(sum_rewards[:n_episodes])), |
| "avg_max_reward": float(np.nanmean(max_rewards[:n_episodes])), |
| "pc_success": float(np.nanmean(all_successes[:n_episodes]) * 100), |
| "eval_s": time.time() - start, |
| "eval_ep_s": (time.time() - start) / n_episodes, |
| + **{ |
| + f"pc_{key}": float(np.mean(all_stage_metrics[key][:n_episodes]) * 100) |
| + for key in STAGE_METRIC_KEYS |
| + }, |
| }, |
| } |
| |
| @@ -774,7 +802,7 @@ def eval_main(cfg: EvalPipelineConfig): |
| env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env, policy_cfg=cfg.policy) |
| |
| recording_dir = Path(cfg.output_dir) / "recordings" if cfg.eval.recording else None |
| - max_episodes_rendered = 0 if cfg.eval.recording else 10 |
| + max_episodes_rendered = 0 if cfg.eval.recording else cfg.eval.max_episodes_rendered |
| videos_dir = None if cfg.eval.recording else Path(cfg.output_dir) / "videos" |
| |
| with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(): |
| @@ -820,9 +848,20 @@ class TaskMetrics(TypedDict): |
| successes: list[bool] |
| video_paths: list[str] |
| predicted_video_paths: list[str] |
| - |
| - |
| -ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths", "predicted_video_paths") |
| + reached_handle: list[bool] |
| + valid_grasp: list[bool] |
| + door_started_moving: list[bool] |
| + grasp_maintained: list[bool] |
| + |
| + |
| +ACC_KEYS = ( |
| + "sum_rewards", |
| + "max_rewards", |
| + "successes", |
| + *STAGE_METRIC_KEYS, |
| + "video_paths", |
| + "predicted_video_paths", |
| +) |
| |
| |
| def eval_one( |
| @@ -872,6 +911,7 @@ def eval_one( |
| successes=[ep["success"] for ep in per_episode], |
| video_paths=task_result.get("video_paths", []), |
| predicted_video_paths=task_result.get("predicted_video_paths", []), |
| + **{key: [ep[key] for ep in per_episode] for key in STAGE_METRIC_KEYS}, |
| ) |
| |
| |
| @@ -990,6 +1030,8 @@ def eval_policy_all( |
| _append("sum_rewards", metrics.get("sum_rewards")) |
| _append("max_rewards", metrics.get("max_rewards")) |
| _append("successes", metrics.get("successes")) |
| + for key in STAGE_METRIC_KEYS: |
| + _append(key, metrics.get(key)) |
| for key in ("video_paths", "predicted_video_paths"): |
| paths = metrics.get(key, []) |
| if paths: |
| @@ -1073,6 +1115,7 @@ def eval_policy_all( |
| "avg_max_reward": _agg_from_list(acc["max_rewards"]), |
| "pc_success": _agg_from_list(acc["successes"]) * 100 if acc["successes"] else float("nan"), |
| "n_episodes": len(acc["sum_rewards"]), |
| + **{f"pc_{key}": _agg_from_list(acc[key]) * 100 for key in STAGE_METRIC_KEYS}, |
| "video_paths": list(acc["video_paths"]), |
| "predicted_video_paths": list(acc["predicted_video_paths"]), |
| } |
| @@ -1085,6 +1128,7 @@ def eval_policy_all( |
| "n_episodes": len(overall["sum_rewards"]), |
| "eval_s": time.time() - start_t, |
| "eval_ep_s": (time.time() - start_t) / max(1, len(overall["sum_rewards"])), |
| + **{f"pc_{key}": _agg_from_list(overall[key]) * 100 for key in STAGE_METRIC_KEYS}, |
| "video_paths": list(overall["video_paths"]), |
| "predicted_video_paths": list(overall["predicted_video_paths"]), |
| } |
| |
| |
| |
| |
| @@ -129,7 +129,16 @@ def make_sample_weighter( |
| # No-op weighter that returns uniform weights |
| return UniformWeighter(device=device) |
| |
| - raise ValueError(f"Unknown sample weighting type: '{config.type}'. Supported types: 'rabc', 'uniform'") |
| + if config.type == "task_balance": |
| + return TaskBalanceWeighter( |
| + device=device, |
| + task_weights=config.extra_params.get("task_weights", {}), |
| + ) |
| + |
| + raise ValueError( |
| + f"Unknown sample weighting type: '{config.type}'. " |
| + "Supported types: 'rabc', 'uniform', 'task_balance'" |
| + ) |
| |
| |
| def _make_rabc_weighter( |
| @@ -237,3 +246,29 @@ class UniformWeighter(SampleWeighter): |
| def get_stats(self) -> dict: |
| """Return empty stats for uniform weighting.""" |
| return {"type": "uniform"} |
| + |
| + |
| +class TaskBalanceWeighter(SampleWeighter): |
| + """Apply explicit inverse-frequency weights using ``batch['task_index']``.""" |
| + |
| + def __init__(self, device: torch.device, task_weights: dict[str, float]): |
| + if not task_weights: |
| + raise ValueError("task_balance requires extra_params.task_weights") |
| + self.device = device |
| + self.task_weights = {int(key): float(value) for key, value in task_weights.items()} |
| + if any(value <= 0 for value in self.task_weights.values()): |
| + raise ValueError("task_balance weights must be positive") |
| + |
| + def compute_batch_weights(self, batch: dict) -> tuple[torch.Tensor, dict]: |
| + if "task_index" not in batch or not isinstance(batch["task_index"], torch.Tensor): |
| + raise KeyError("task_balance requires a tensor batch['task_index']") |
| + task_indices = batch["task_index"].reshape(-1).to(self.device, dtype=torch.long) |
| + weights = torch.empty(task_indices.shape[0], device=self.device, dtype=torch.float32) |
| + for task_index in task_indices.unique().tolist(): |
| + if task_index not in self.task_weights: |
| + raise KeyError(f"No task_balance weight configured for task_index={task_index}") |
| + weights[task_indices == task_index] = self.task_weights[task_index] |
| + return weights, {"mean_weight": weights.mean().item(), "type": "task_balance"} |
| + |
| + def get_stats(self) -> dict: |
| + return {"type": "task_balance", "task_weights": self.task_weights} |
|
|
