From e50308789c8c21bec0c256349e126424aba02116 Mon Sep 17 00:00:00 2001 From: Gangwei XU Date: Mon, 15 Jun 2026 15:52:48 +0800 Subject: [PATCH] fix(lingbot-va): align RoboTwin evaluation (#3784) Thank you for the RoboTwin fix, and alignment! --- src/lerobot/envs/configs.py | 2 +- src/lerobot/envs/robotwin.py | 87 +++++++++++++++++- .../lingbot_va/modeling_lingbot_va.py | 16 +++- src/lerobot/scripts/lerobot_eval.py | 88 ++++++++++++------- 4 files changed, 157 insertions(+), 36 deletions(-) diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index f7f0b01a5..3624357e2 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -757,7 +757,7 @@ class RoboTwinEnvConfig(EnvConfig): task: str = "beat_block_hammer" # single task or comma-separated list fps: int = 25 - episode_length: int = 300 + episode_length: int = 1200 obs_type: str = "pixels_agent_pos" render_mode: str = "rgb_array" # Available cameras from RoboTwin's aloha-agilex embodiment: head_camera diff --git a/src/lerobot/envs/robotwin.py b/src/lerobot/envs/robotwin.py index f5ea01999..16471c5d9 100644 --- a/src/lerobot/envs/robotwin.py +++ b/src/lerobot/envs/robotwin.py @@ -17,6 +17,7 @@ from __future__ import annotations import importlib import logging +import os from collections import defaultdict from collections.abc import Callable, Sequence from functools import partial @@ -47,7 +48,10 @@ ACTION_DIM = 14 # 7 DOF × 2 arms (joint-space control mode) EEF_ACTION_DIM = 16 ACTION_LOW = -1.0 ACTION_HIGH = 1.0 -DEFAULT_EPISODE_LENGTH = 300 +DEFAULT_EPISODE_LENGTH = 1200 +OFFICIAL_INSTRUCTION_ENV = "LEROBOT_ROBOTWIN_OFFICIAL_INSTRUCTION" +OFFICIAL_INSTRUCTION_TYPE_ENV = "LEROBOT_ROBOTWIN_INSTRUCTION_TYPE" +OFFICIAL_INSTRUCTION_MAX_ENV = "LEROBOT_ROBOTWIN_INSTRUCTION_MAX" def _compose_eef_pose(new_pose: np.ndarray, init_pose: np.ndarray) -> np.ndarray: @@ -77,6 +81,78 @@ def _add_init_eef_pose(delta_pose: np.ndarray, init_pose: np.ndarray) -> np.ndar return out +def _env_flag(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() in {"1", "true", "yes", "on"} + + +def _arm_for_block(block: Any) -> str: + return "left" if float(block.get_pose().p[0]) < 0 else "right" + + +def _robotwin_blocks_episode_info(task_name: str, env: Any) -> dict[str, str] | None: + """Infer the episode-info dict used by RoboTwin's official instruction generator for block ranking.""" + if task_name == "blocks_ranking_rgb": + return { + "{A}": "red block", + "{B}": "green block", + "{C}": "blue block", + "{a}": _arm_for_block(env.block1), + "{b}": _arm_for_block(env.block2), + "{c}": _arm_for_block(env.block3), + } + if task_name == "blocks_ranking_size": + return { + "{A}": "large block", + "{B}": "medium block", + "{C}": "small block", + "{a}": _arm_for_block(env.block1), + "{b}": _arm_for_block(env.block2), + "{c}": _arm_for_block(env.block3), + } + return None + + +def _generate_robotwin_official_instruction(task_name: str, env: Any) -> str: + """Generate language with RoboTwin's official task templates, matching its eval client.""" + fallback = task_name.replace("_", " ") + episode_info = _robotwin_blocks_episode_info(task_name, env) + if episode_info is None: + logger.warning("Official RoboTwin instruction is not implemented for task=%s; using %r.", task_name, fallback) + return fallback + + try: + from description.utils.generate_episode_instructions import generate_episode_descriptions + except Exception: + logger.warning("Failed to import RoboTwin official instruction generator; using %r.", fallback, exc_info=True) + return fallback + + instruction_type = os.environ.get(OFFICIAL_INSTRUCTION_TYPE_ENV, "seen") + try: + max_descriptions = int(os.environ.get(OFFICIAL_INSTRUCTION_MAX_ENV, "1000000")) + except ValueError: + max_descriptions = 1000000 + + results = generate_episode_descriptions(task_name, [episode_info], max_descriptions=max_descriptions) + if not results: + logger.warning("RoboTwin generated no official instructions for task=%s; using %r.", task_name, fallback) + return fallback + + options = results[0].get(instruction_type) or results[0].get("seen") or results[0].get("unseen") + if not options: + logger.warning( + "RoboTwin generated no %s official instructions for task=%s; using %r.", + instruction_type, + task_name, + fallback, + ) + return fallback + + return str(np.random.choice(options)) + + # D435 dims from task_config/_camera_config.yml (what demo_clean.yml selects). DEFAULT_CAMERA_H = 240 DEFAULT_CAMERA_W = 320 @@ -382,6 +458,15 @@ class RoboTwinEnv(gym.Env): self.episode_index += self._reset_stride self._step_count = 0 + use_official_instruction = self.task_name in {"blocks_ranking_rgb", "blocks_ranking_size"} + if _env_flag(OFFICIAL_INSTRUCTION_ENV, default=use_official_instruction): + self.task_description = _generate_robotwin_official_instruction(self.task_name, self._env) + if hasattr(self._env, "set_instruction"): + self._env.set_instruction(instruction=self.task_description) + logger.info("RoboTwin official instruction | task=%s | %s", self.task_name, self.task_description) + else: + self.task_description = self.task_name.replace("_", " ") + # In eef mode the policy predicts pose deltas relative to the initial eef pose. if self.action_mode == "ee": self._init_eef_pose = self._read_eef_pose() diff --git a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py index b1040283c..e3446eebf 100644 --- a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py @@ -1175,6 +1175,7 @@ class LingBotVAPolicy(PreTrainedPolicy): self._frozen: dict = {} self.last_predicted_frames: Tensor | None = None + self.last_predicted_latents: Tensor | None = None self.reset() # Frozen-module lazy loading (VAE + UMT5 + tokenizer) @@ -1245,6 +1246,7 @@ class LingBotVAPolicy(PreTrainedPolicy): self._prompt_embeds = None self._negative_prompt_embeds = None self.last_predicted_frames = None + self.last_predicted_latents = None self._use_cfg = (cfg.guidance_scale > 1) or (cfg.action_guidance_scale > 1) # Two independent flow-matching schedulers (video latent + action streams). self._scheduler = FlowMatchScheduler(shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True) @@ -1533,7 +1535,10 @@ class LingBotVAPolicy(PreTrainedPolicy): self._executed_actions = actions if self.config.save_predicted_video: - self.last_predicted_frames = self._decode_predicted_video(latents) + # Match upstream LingBot-VA visualization: collect chunk latents and decode the + # concatenated latent sequence once after the rollout finishes. + self.last_predicted_frames = None + self.last_predicted_latents = latents.detach().to("cpu") # On the first chunk, frame 0 is the conditioning frame (already "known"): the upstream # LIBERO client skips it (start_idx=1), so we drop the first frame's actions here. @@ -1908,13 +1913,20 @@ class LingBotVAPolicy(PreTrainedPolicy): return actions, latents # Predicted-video decoding (opt-in) + @torch.no_grad() + def decode_predicted_latents(self, latents) -> Tensor: + """Decode a concatenated predicted-latent sequence into ``[T, H, W, 3]`` uint8 frames.""" + return self._decode_predicted_video(latents) + @torch.no_grad() def _decode_predicted_video(self, latents) -> Tensor: """VAE-decode predicted latents into a uint8 frame stack ``[T, H, W, 3]`` on CPU.""" vae = self._vae z_dim = vae.config.z_dim + vae_device = next(vae.parameters()).device + latents = latents.to(device=vae_device, dtype=vae.dtype) latents = denormalize_latents( - latents.to(vae.dtype), vae.config.latents_mean, vae.config.latents_std, z_dim + latents, vae.config.latents_mean, vae.config.latents_std, z_dim ) video = vae.decode(latents, return_dict=False)[0] # [B, C, F, H, W] in [-1, 1] video = (video.float().clamp(-1, 1) + 1.0) / 2.0 diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 406a12373..a39e00ce5 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -169,7 +169,7 @@ def rollout( env_features: dict | None = None, recording_repo_id: str | None = None, recording_private: bool = False, - predicted_frames_callback: Callable[[PreTrainedPolicy], None] | None = None, + predicted_latents_callback: Callable[[PreTrainedPolicy], None] | None = None, ) -> dict: """Run a batched policy rollout once through a batch of environments. @@ -199,9 +199,9 @@ def rollout( are returned optionally because they typically take more memory to cache. Defaults to False. render_callback: Optional rendering callback to be used after the environments are reset, and after every step. - predicted_frames_callback: Optional callback invoked after every ``select_action`` with the policy - itself. World-model policies (e.g. LingBot-VA) stash their decoded predicted video frames on - ``policy.last_predicted_frames``; this lets the caller collect them to save predicted-video MP4s. + predicted_latents_callback: Optional callback invoked after every ``select_action`` with the policy + itself. World-model policies (e.g. LingBot-VA) stash predicted video latents on + ``policy.last_predicted_latents``; this lets the caller concatenate chunks and decode once. Returns: The dictionary described above. """ @@ -280,8 +280,8 @@ def rollout( observation = preprocessor(observation) with torch.inference_mode(): action = policy.select_action(observation) - if predicted_frames_callback is not None: - predicted_frames_callback(policy) + if predicted_latents_callback is not None: + predicted_latents_callback(policy) action = postprocessor(action) action_transition = {ACTION: action} @@ -301,18 +301,26 @@ def rollout( # available if none of the envs finished. if "final_info" in info: final_info = info["final_info"] - if not isinstance(final_info, dict): - raise RuntimeError( - "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " - "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + if isinstance(final_info, dict): + is_success = final_info.get("is_success", [False] * env.num_envs) + successes = ( + is_success.tolist() + if hasattr(is_success, "tolist") + else [bool(is_success)] * env.num_envs ) - successes = final_info["is_success"].tolist() + else: + # Gymnasium < 1.0 returns final_info as a per-env sequence/object array, + # with entries set to a dict only for envs that just finished. + successes = [] + for item in final_info: + if isinstance(item, dict) and "is_success" in item: + successes.append(bool(item["is_success"])) + else: + successes.append(False) elif "is_success" in info: is_success = info["is_success"] successes = ( - is_success.tolist() - if hasattr(is_success, "tolist") - else [bool(is_success)] * env.num_envs + is_success.tolist() if hasattr(is_success, "tolist") else [bool(is_success)] * env.num_envs ) else: successes = [False] * env.num_envs @@ -479,14 +487,15 @@ def eval_policy( predicted_video_paths: list[str] = [] n_predicted_rendered = 0 - # Collects the policy's decoded predicted-video frames across a rollout (world-model policies only). - def collect_predicted_frames(policy: PreTrainedPolicy): - frames = getattr(policy, "last_predicted_frames", None) - if frames is not None: - pred_frames.append( - np.asarray(frames.detach().to("cpu")) if hasattr(frames, "detach") else np.asarray(frames) + # Collect predicted-video latents across a rollout (world-model policies only). The latents are + # concatenated and decoded once after the rollout, matching upstream LingBot-VA's visualization path. + def collect_predicted_latents(policy: PreTrainedPolicy): + latents = getattr(policy, "last_predicted_latents", None) + if latents is not None: + pred_latents.append( + latents.detach().to("cpu") if hasattr(latents, "detach") else torch.as_tensor(latents).cpu() ) - policy.last_predicted_frames = None + policy.last_predicted_latents = None if return_episode_data: episode_data: dict | None = None @@ -500,7 +509,7 @@ def eval_policy( ep_frames: list[np.ndarray] = [] if save_predicted_video: - pred_frames: list[np.ndarray] = [] + pred_latents: list[torch.Tensor] = [] if start_seed is None: seeds = None @@ -522,7 +531,7 @@ def eval_policy( env_features=env_features, recording_repo_id=recording_repo_id, recording_private=recording_private, - predicted_frames_callback=collect_predicted_frames if save_predicted_video else None, + predicted_latents_callback=collect_predicted_latents if save_predicted_video else None, ) # Figure out where in each rollout sequence the first done condition was encountered (results after @@ -589,9 +598,19 @@ def eval_policy( n_episodes_rendered += 1 # Maybe save the policy's predicted (imagined) video for this batch's rollout. - if save_predicted_video and len(pred_frames) > 0: - # pred_frames is a list of [F, H, W, C] uint8 stacks emitted on chunk refills; concat over time. - predicted_video = np.concatenate(pred_frames, axis=0) + if save_predicted_video and len(pred_latents) > 0: + predicted_latent = torch.cat(pred_latents, dim=2) + decoder = getattr(policy, "decode_predicted_latents", None) or getattr( + policy, "_decode_predicted_video", None + ) + if decoder is None: + raise AttributeError( + "Policy config requested predicted-video saving, but the policy does not expose " + "`decode_predicted_latents` or `_decode_predicted_video`." + ) + predicted_video = decoder(predicted_latent) + if hasattr(predicted_video, "detach"): + predicted_video = predicted_video.detach().to("cpu").numpy() videos_dir.mkdir(parents=True, exist_ok=True) predicted_video_path = videos_dir / f"pred_episode_{n_predicted_rendered}.mp4" predicted_video_paths.append(str(predicted_video_path)) @@ -793,9 +812,10 @@ class TaskMetrics(TypedDict): max_rewards: list[float] successes: list[bool] video_paths: list[str] + predicted_video_paths: list[str] -ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths") +ACC_KEYS = ("sum_rewards", "max_rewards", "successes", "video_paths", "predicted_video_paths") def eval_one( @@ -844,6 +864,7 @@ def eval_one( max_rewards=[ep["max_reward"] for ep in per_episode], successes=[ep["success"] for ep in per_episode], video_paths=task_result.get("video_paths", []), + predicted_video_paths=task_result.get("predicted_video_paths", []), ) @@ -904,6 +925,7 @@ def run_one( if max_episodes_rendered > 0: metrics.setdefault("video_paths", []) + metrics.setdefault("predicted_video_paths", []) return task_group, task_id, metrics @@ -961,11 +983,11 @@ def eval_policy_all( _append("sum_rewards", metrics.get("sum_rewards")) _append("max_rewards", metrics.get("max_rewards")) _append("successes", metrics.get("successes")) - # video_paths is list-like - paths = metrics.get("video_paths", []) - if paths: - group_acc[group]["video_paths"].extend(paths) - overall["video_paths"].extend(paths) + for key in ("video_paths", "predicted_video_paths"): + paths = metrics.get(key, []) + if paths: + group_acc[group][key].extend(paths) + overall[key].extend(paths) # Choose runner (sequential vs threaded) task_runner = partial( @@ -1037,6 +1059,7 @@ def eval_policy_all( "pc_success": _agg_from_list(acc["successes"]) * 100 if acc["successes"] else float("nan"), "n_episodes": len(acc["sum_rewards"]), "video_paths": list(acc["video_paths"]), + "predicted_video_paths": list(acc["predicted_video_paths"]), } # overall aggregates @@ -1048,6 +1071,7 @@ def eval_policy_all( "eval_s": time.time() - start_t, "eval_ep_s": (time.time() - start_t) / max(1, len(overall["sum_rewards"])), "video_paths": list(overall["video_paths"]), + "predicted_video_paths": list(overall["predicted_video_paths"]), } return {