VLA/wall_x/infer/env_libero.py

import os
import json
import numpy as np
from typing import Dict, Any, Tuple, List
from libero.libero import benchmark
from wall_x.infer.env import BaseEnv, InferConfig
from wall_x.serving.policy.wall_x_policy import WallXPolicy

from wall_x.infer.base_dataclass import RobotStateActionData
from wall_x.infer.utils_libero import (
    get_libero_env,
    get_libero_dummy_action,
    get_libero_image,
    get_libero_wrist_image,
    quat2axisangle,
    TaskSuite,
    save_rollout_video,
)
from robosuite.wrappers import VisualizationWrapper


def _create_libero_env_standalone(
    task_id: int,
    task_suite_name: str,
    model_family: str = "wallx",
    resolution: int = 256,
    seed: int = 7,
) -> Any:
    """
    Standalone function to create a Libero environment, independent of LiberoRobotEnv instance.
    Used for creating environments in subprocess during multi-batch inference,
    avoiding serialization of large objects containing the model.

    Args:
        task_id: Task ID
        task_suite_name: Task suite name
        model_family: Model family
        resolution: Resolution
        seed: Random seed

    Returns:
        Environment instance
    """
    from libero.libero import benchmark

    # Get task suite and task
    benchmark_dict = benchmark.get_benchmark_dict()
    task_suite = benchmark_dict[task_suite_name]()
    task = task_suite.get_task(task_id)

    # Create environment
    env, _ = get_libero_env(
        task,
        model_family=model_family,
        resolution=resolution,
        seed=seed,
    )

    # Wrap environment
    env.env = VisualizationWrapper(env.env)
    env.env.set_visualization_setting(setting="grippers", visible=False)

    return env


class LiberoRobotEnv(BaseEnv):
    def __init__(
        self,
        config: InferConfig,
        task_suite_name: str = TaskSuite.LIBERO_SPATIAL,
        initial_states_path: str = "DEFAULT",
        rollout_dir: str = "./rollouts",
        model_family: str = "wallx",
        resolution: int = 256,
        seed: int = 7,
    ):

        super().__init__(config)
        self.logger.info(
            f"Initializing LiberoRobotEnv (Stateless), task suite: {task_suite_name}"
        )

        self.model = self._register_model()
        self.model_family = model_family
        self.resolution = resolution
        self.seed = seed

        self.logger.info("Importing Libero and related utils...")
        self.RobotStateActionData = RobotStateActionData

        self.rollout_dir = os.path.join(rollout_dir, task_suite_name)
        os.makedirs(self.rollout_dir, exist_ok=True)

        if save_rollout_video is not None:
            self.save_rollout_video = save_rollout_video
        else:
            self.save_rollout_video = None
            self.logger.warning("save_rollout_video not found, video saving disabled.")

        self.task_suite_name = task_suite_name
        benchmark_dict = benchmark.get_benchmark_dict()
        self.task_suite = benchmark_dict[self.task_suite_name]()
        self.num_tasks = self.task_suite.n_tasks

        self.initial_states_path = initial_states_path
        self.all_initial_states = None
        if self.initial_states_path != "DEFAULT":
            try:
                with open(self.initial_states_path, "r") as f:
                    self.all_initial_states = json.load(f)
                self.logger.info(
                    f"Loaded custom initial states from {self.initial_states_path}"
                )
            except Exception as e:
                self.logger.error(f"Failed to load initial states file: {e}")
                raise

    def _register_model(self) -> WallXPolicy:

        return WallXPolicy(
            model_path=self.config.model_path,
            train_config=self.config.train_config,
            action_tokenizer_path=self.config.action_tokenizer_path,
            action_dim=self.config.action_dim,
            agent_pos_dim=self.config.action_dim,
            pred_horizon=self.config.pred_horizon,
            camera_key=self.config.cam_names,
            predict_mode=self.config.predict_mode,
        )

    def get_instruction(self, task_desc: str) -> str:
        return task_desc

    def get_observation(self, raw_obs: Dict[str, Any]) -> Dict[str, Any]:
        if raw_obs is None:
            raise ValueError("Raw observation is None")

        data_obj = self.RobotStateActionData(config=self.config)

        pos = raw_obs["robot0_eef_pos"]
        rot = quat2axisangle(raw_obs["robot0_eef_quat"])
        grip = raw_obs["robot0_gripper_qpos"][0:1]

        data_obj.save_state_data_with_key(pos[None], "follow_right_ee_cartesian_pos")
        data_obj.save_state_data_with_key(rot[None], "follow_right_ee_rotation")
        data_obj.save_state_data_with_key(grip[None], "follow_right_gripper")
        data_obj.dof_mask = self._get_dof_mask()

        face_view = get_libero_image(raw_obs)
        right_wrist_view = get_libero_wrist_image(raw_obs)

        return {
            "robot_state_action_data": data_obj,
            "face_view": face_view,
            "right_wrist_view": right_wrist_view,
        }

    def apply_action(
        self, input_data: Dict[str, Any], env: Any = None, replay_images: list = None
    ) -> bool:
        if env is None:
            raise ValueError(
                "In Stateless mode, apply_action must be called with explicit 'env' parameter"
            )

        action_data = input_data["robot_state_action_data"]
        right_arm_traj = self._get_right_arm_action(action_data)
        while (
            right_arm_traj is not None
            and right_arm_traj.ndim > 2
            and right_arm_traj.shape[0] == 1
        ):
            right_arm_traj = right_arm_traj.squeeze(0)

        done = False
        t = 0

        try:
            for i in range(len(right_arm_traj)):
                if done:
                    break

                action_7d = right_arm_traj[i]
                obs, reward, done, info = env.step(action_7d)
                t += 1

                if obs is not None and replay_images is not None:
                    replay_images.append(get_libero_image(obs))

                input_data["_last_obs"] = obs

        except Exception as e:
            self.logger.error(f"Env step error: {e}")
            return False  # Error treated as failure

        return done, t

    def apply_action_batch(
        self,
        vec_env: Any,
        trajectories: List[np.ndarray],
        active_indices: List[int],
        status_list: List[Dict[str, Any]],
        model_outputs: List[Dict[str, Any]],
    ) -> None:
        """
        Execute action trajectories in parallel batch.

        Uses SubprocVectorEnv to execute actions in parallel for all active environments.
        Integrates vec_env.step(batch_actions, id=still_active) in this function.
        """
        if not trajectories:
            return

        max_traj_len = max(len(traj) for traj in trajectories)
        if max_traj_len == 0:
            return

        for step_idx in range(max_traj_len):
            # Check if there are still active environments
            still_active = [
                idx
                for idx in active_indices
                if (not status_list[idx]["done"])
                and status_list[idx]["count"] > 0
                and step_idx < len(trajectories[active_indices.index(idx)])
            ]
            if not still_active:
                break

            # Build batch actions (only includes actions for still_active environments)
            batch_actions = []
            for idx in still_active:
                traj_idx = active_indices.index(idx)
                action_7d = trajectories[traj_idx][step_idx]
                # Ensure action_7d is numpy array or list
                if isinstance(action_7d, np.ndarray):
                    batch_actions.append(action_7d)
                else:
                    batch_actions.append(np.array(action_7d))

            # Convert to numpy array with shape (batch_size, action_dim)
            batch_actions = np.array(batch_actions)

            # Execute step in parallel (only for still_active environments)
            obs_list, reward_list, done_list, info_list = vec_env.step(
                batch_actions, id=still_active
            )

            # Process returned results
            if obs_list.dtype == object:
                obs_list = list(obs_list)
            else:
                obs_list = [obs_list[i] for i in range(len(obs_list))]
            done_list = [bool(done_list[i]) for i in range(len(done_list))]

            # Update each environment's status
            for i, idx in enumerate(still_active):
                st = status_list[idx]
                obs = obs_list[i]
                done = done_list[i]

                if obs is not None:
                    st["current_obs"] = obs
                    if st["replay_images"] is not None:
                        st["replay_images"].append(get_libero_image(obs))

                st["count"] -= 1
                st["success"] = done
                st["done"] = done or st["count"] <= 0

                # Update model_output's _last_obs
                model_outputs[active_indices.index(idx)]["_last_obs"] = obs

    def get_task_info(self, task_id: int) -> Tuple[str, Any]:
        """
        Get task information (task description and initial states) without creating environment.

        Returns:
            Tuple[str, Any]: (task_desc, default_initial_states)
        """
        if task_id < 0 or task_id >= self.num_tasks:
            raise ValueError(f"Invalid task ID: {task_id}")

        task = self.task_suite.get_task(task_id)
        task_desc = task.language
        default_initial_states = self.task_suite.get_task_init_states(task_id)

        return task_desc, default_initial_states

    def create_env_for_task(self, task_id: int) -> Tuple[Any, str, Any]:
        if task_id < 0 or task_id >= self.num_tasks:
            raise ValueError(f"Invalid task ID: {task_id}")

        task = self.task_suite.get_task(task_id)
        default_initial_states = self.task_suite.get_task_init_states(task_id)

        env, task_desc = get_libero_env(
            task,
            model_family=self.model_family,
            resolution=self.resolution,
            seed=self.seed,
        )

        env.env = VisualizationWrapper(env.env)
        env.env.set_visualization_setting(setting="grippers", visible=False)

        return env, task_desc, default_initial_states

    def _get_initial_state_for_episode(
        self, task_desc: str, default_states: Any, episode_idx: int
    ) -> np.ndarray:
        if self.initial_states_path == "DEFAULT":
            if default_states is None:
                raise ValueError("Default states missing")
            return default_states[episode_idx]
        else:
            if self.all_initial_states is None:
                raise ValueError("Custom states not loaded")
            initial_states_task_key = task_desc.replace(" ", "_")
            episode_key = f"demo_{episode_idx}"
            if not self.all_initial_states[initial_states_task_key][episode_key][
                "success"
            ]:
                raise ValueError(f"Expert demo failed for {episode_key}")
            return np.array(
                self.all_initial_states[initial_states_task_key][episode_key][
                    "initial_state"
                ]
            )

    def reset_env(
        self, env: Any, task_desc: str, default_states: Any, episode_idx: int
    ) -> Any:
        try:
            if episode_idx >= 0:
                state = self._get_initial_state_for_episode(
                    task_desc, default_states, episode_idx
                )
                obs = env.set_init_state(state)
                return obs
            else:
                return env.reset()
        except Exception as e:
            self.logger.error(f"Reset failed: {e}, falling back to default reset")
            return env.reset()

    def _get_dof_mask(self):
        dof_config = self.config.train_config["dof_config"]
        total_dof = sum(dof_config.values())
        dof_mask = np.ones((1, self.config.action_horizon, total_dof))
        mask_keys = [
            "follow_left_ee_cartesian_pos",
            "follow_left_ee_rotation",
            "follow_left_gripper",
            "head_actions",
            "height",
            "velocity_decomposed",
        ]
        start_idx = 0
        for key, dof_size in dof_config.items():
            if key in mask_keys:
                dof_mask[:, :, start_idx : start_idx + dof_size] = 0
            start_idx += dof_size
        return dof_mask

    def _get_right_arm_action(
        self, robot_state_action_data: RobotStateActionData
    ) -> np.ndarray:
        right_ee_cartesian_pos = robot_state_action_data.data[
            "action_right_ee_cartesian_pos"
        ]
        right_ee_rotation = robot_state_action_data.data["action_right_ee_rotation"]
        right_gripper = robot_state_action_data.data["action_right_gripper"]
        return np.concatenate(
            [right_ee_cartesian_pos, right_ee_rotation, right_gripper], axis=1
        )

    def _get_left_arm_action(
        self, robot_state_action_data: RobotStateActionData
    ) -> np.ndarray:
        left_ee_cartesian_pos = robot_state_action_data.data[
            "action_left_ee_cartesian_pos"
        ]
        left_ee_rotation = robot_state_action_data.data["action_left_ee_rotation"]
        left_gripper = robot_state_action_data.data["action_left_gripper"]
        return np.concatenate(
            [left_ee_cartesian_pos, left_ee_rotation, left_gripper], axis=1
        )

    def _save_rollout(
        self,
        replay_images: List[np.ndarray],
        success: bool,
        task_id: int,
        task_desc: str,
        episode_idx: int,
    ):
        if not self.save_rollout_video or not replay_images:
            return
        try:
            task_name_safe = task_desc.replace(" ", "_").replace(".", "")
            filename = f"{episode_idx}{'_SUCCESS' if success else '_FAILURE'}--_{task_name_safe}.mp4"
            self.save_rollout_video(
                self.rollout_dir,
                replay_images,
                filename,
                success=success,
                task_description=task_desc,
                log_file=None,
                model_family=self.model_family,
            )
            self.logger.info(f"Saved video: {filename}")
        except Exception as e:
            self.logger.error(f"Save video failed: {e}")

    def run_infer_flow_action(
        self,
        env: Any,
        task_id: int,
        task_desc: str,
        default_initial_states: Any,
        episode_idx: int,
        max_infer_times: int = 5,
        num_steps_wait: int = 10,
    ) -> bool:
        replay_images = []
        num_steps = 0
        done = False
        count = max_infer_times

        current_obs = self.reset_env(
            env, task_desc, default_initial_states, episode_idx
        )
        if current_obs is None:
            return False

        while num_steps < num_steps_wait:
            obs, reward, done, info = env.step(
                get_libero_dummy_action(self.model_family)
            )
            num_steps += 1
            if obs is not None:
                current_obs = obs

        while not done and count > 0:
            try:
                model_input = self.get_observation(current_obs)
                instruction = self.get_instruction(task_desc)
                model_input["prompt"] = instruction
                model_input["dataset_names"] = "libero_all"

                state = np.concatenate(
                    [
                        model_input["robot_state_action_data"].data[
                            "state_right_ee_cartesian_pos"
                        ],
                        model_input["robot_state_action_data"].data[
                            "state_right_ee_rotation"
                        ],
                        model_input["robot_state_action_data"].data[
                            "state_right_gripper"
                        ],
                    ],
                    axis=-1,
                )

                model_input["state"] = state
                model_output = self.model.infer(model_input)

                model_output["robot_state_action_data"] = model_input[
                    "robot_state_action_data"
                ]
                model_output["robot_state_action_data"].save_action_data(
                    model_output["predict_action"]
                )

                model_output["_last_obs"] = None

                done, delta_t = self.apply_action(
                    model_output, env=env, replay_images=replay_images
                )

                if model_output.get("_last_obs") is not None:
                    current_obs = model_output["_last_obs"]

                count -= delta_t

            except Exception as e:
                self.logger.error(f"Episode Error: {e}")
                import traceback

                traceback.print_exc()
                break

        success = done
        if count <= 0 and not done:
            self.logger.warning(
                f"Timeout: reached {max_infer_times} steps without success."
            )
            success = False

        self._save_rollout(replay_images, success, task_id, task_desc, episode_idx)
        return success

    def run_infer_flow_action_batch(
        self,
        vec_env: Any,
        task_ids: List[int] = None,
        task_descs: List[str] = None,
        default_initial_states_list: List[Any] = None,
        episode_indices: List[int] = None,
        max_infer_times: int = 5,
        num_steps_wait: int = 10,
    ) -> List[bool]:
        """
        Support batch inference: model inference in parallel (batch), environment execution in parallel (SubprocVectorEnv).

        Uses SubprocVectorEnv to run environments in subprocess during multi-batch inference, all environments execute actions in parallel.

        Returns a list of success flags for each sample.
        """
        if vec_env is None:
            raise ValueError("vec_env must be specified")
        batch_size = len(vec_env)
        if task_ids is not None:
            assert len(task_ids) == batch_size, "task_ids length must match envs"
        if episode_indices is not None:
            assert (
                len(episode_indices) == batch_size
            ), "episode_indices length must match envs"

        status_list = []
        for i in range(batch_size):
            status_list.append(
                {
                    "vec_env": vec_env,
                    "env_id": i,  # Index in vec_env
                    "task_desc": task_descs[i],
                    "replay_images": [],
                    "num_steps": 0,
                    "done": False,  # Whether episode has ended
                    "success": False,  # Whether successfully completed
                    "count": max_infer_times,
                    "current_obs": None,
                    "default_states": None,
                }
            )

        # Initialize/reset: Use SubprocVectorEnv to batch set initial states
        init_states_to_set = []
        for i in range(batch_size):
            task_desc = task_descs[i]
            default_states = default_initial_states_list[i]
            status_list[i]["default_states"] = default_states
            ep_i = episode_indices[i]
            init_state = self._get_initial_state_for_episode(
                task_desc, default_states, ep_i
            )
            init_states_to_set.append(init_state)

        # Batch set initial states
        try:
            obs_list = vec_env.set_init_state(init_states_to_set)
            if obs_list.dtype == object:
                obs_list = list(obs_list)
            else:
                obs_list = [obs_list[i] for i in range(len(obs_list))]

            for i, obs in enumerate(obs_list):
                if obs is None:
                    raise ValueError(
                        f"Reset environment returned None, task_id: {task_ids[i]}, episode_idx: {episode_indices[i]}"
                    )
                status_list[i]["current_obs"] = obs
                status_list[i]["done"] = False
                status_list[i]["success"] = False
                status_list[i]["count"] = max_infer_times
        except Exception as e:
            self.logger.error(f"Failed to batch set initial states: {e}")
            raise

        # Warmup steps (batch execution)
        dummy_action = get_libero_dummy_action(self.model_family)
        dummy_actions = np.array([dummy_action] * batch_size)
        for _ in range(num_steps_wait):
            obs_list, _, done_list, _ = vec_env.step(dummy_actions)
            # Update current_obs
            if obs_list.dtype == object:
                obs_list = list(obs_list)
            else:
                obs_list = [obs_list[i] for i in range(len(obs_list))]
            for i, obs in enumerate(obs_list):
                if obs is not None:
                    status_list[i]["current_obs"] = obs

        # Main loop: model parallel inference, environment parallel execution (SubprocVectorEnv)
        while any((not st["done"]) and st["count"] > 0 for st in status_list):
            active_indices = [
                idx
                for idx, st in enumerate(status_list)
                if (not st["done"]) and st["count"] > 0
            ]
            if not active_indices:
                break
            print(f"Batch infer loop, active indices: {active_indices}")

            observations = []
            instructions = []
            for idx in active_indices:
                st = status_list[idx]
                observations.append(self.get_observation(st["current_obs"]))
                instructions.append(self.get_instruction(st["task_desc"]))

            # Model batch inference
            model_outputs = self.model.infer_flow_action_batch(
                observations, instructions
            )
            # Extract action trajectories for all active environments
            trajectories = []
            for out in model_outputs:
                action_data = out["robot_state_action_data"]
                right_arm_traj = self._get_right_arm_action(action_data)
                while (
                    right_arm_traj is not None
                    and right_arm_traj.ndim > 2
                    and right_arm_traj.shape[0] == 1
                ):
                    right_arm_traj = right_arm_traj.squeeze(0)
                if right_arm_traj is None or len(right_arm_traj) == 0:
                    # If trajectory is empty, create an empty trajectory
                    right_arm_traj = np.array([]).reshape(0, 7)
                trajectories.append(right_arm_traj)

            # Use apply_action_batch to execute action trajectories in parallel
            try:
                self.apply_action_batch(
                    vec_env=vec_env,
                    trajectories=trajectories,
                    active_indices=active_indices,
                    status_list=status_list,
                    model_outputs=model_outputs,
                )
            except Exception as e:
                self.logger.error(f"Batch parallel action error: {e}")
                # Mark all active environments as failed
                for idx in active_indices:
                    status_list[idx]["done"] = True
                    status_list[idx]["success"] = False

        # Save replay and results
        success_list = []
        for i, st in enumerate(status_list):
            success = st.get("success", False)
            if st["count"] <= 0 and not st["success"]:
                self.logger.warning(
                    f"Batch timeout: reached {max_infer_times} steps without success (idx {i})."
                )
            st["replay_images"] = st.get("replay_images", [])
            tid_i = task_ids[i]
            epi_i = episode_indices[i]
            self._save_rollout(
                st["replay_images"],
                success,
                tid_i,
                st["task_desc"],
                epi_i,
            )
            success_list.append(success)

        return success_list

    def run_infer_ar_action(
        self,
        env: Any,
        task_id: int,
        task_desc: str,
        default_initial_states: Any,
        episode_idx: int,
        max_infer_times: int = 10,
        num_steps_wait: int = 10,
    ) -> bool:
        replay_images = []
        num_steps = 0
        done = False
        count = max_infer_times

        current_obs = self.reset_env(
            env, task_desc, default_initial_states, episode_idx
        )
        if current_obs is None:
            self.logger.error("Environment reset returned None.")
            return False

        while num_steps < num_steps_wait:
            obs, reward, done, info = env.step(
                get_libero_dummy_action(self.model_family)
            )
            num_steps += 1
            if obs is not None:
                current_obs = obs

        while not done and count > 0:
            try:
                model_input = self.get_observation(current_obs)
                instruction = self.get_instruction(task_desc)

                model_output = self.model.infer_ar_action(model_input, instruction)

                model_output["_last_obs"] = None

                done, delta_t = self.apply_action(
                    model_output, env=env, replay_images=replay_images
                )

                if model_output.get("_last_obs") is not None:
                    current_obs = model_output["_last_obs"]
                else:
                    if not done:
                        self.logger.warning(
                            "Did not receive new observation after apply_action, but episode is not done."
                        )

                count -= delta_t

            except Exception as e:
                self.logger.error(
                    f"AR Episode Run Error (Task {task_id}, Ep {episode_idx}): {e}"
                )
                import traceback

                traceback.print_exc()
                break

        success = done
        if count <= 0 and not done:
            self.logger.warning(
                f"Timeout: AR policy reached {max_infer_times} steps without success."
            )
            success = False

        self._save_rollout(replay_images, success, task_id, task_desc, episode_idx)

        return success