embodied_gen/envs/pick_embodiedgen.py

# Project EmbodiedGen
#
# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.

import json
import os

import numpy as np
import sapien
import torch
import torchvision.transforms as transforms
from mani_skill.envs.sapien_env import BaseEnv
from mani_skill.sensors.camera import CameraConfig
from mani_skill.utils import sapien_utils
from mani_skill.utils.building import actors
from mani_skill.utils.building.ground import build_ground
from mani_skill.utils.registration import register_env
from mani_skill.utils.structs.actor import Actor
from mani_skill.utils.structs.pose import Pose
from mani_skill.utils.structs.types import (
    GPUMemoryConfig,
    SceneConfig,
    SimConfig,
)
from mani_skill.utils.visualization.misc import tile_images
from tqdm import tqdm
from embodied_gen.models.gs_model import GaussianOperator
from embodied_gen.utils.enum import LayoutInfo, Scene3DItemEnum
from embodied_gen.utils.geometry import bfs_placement, quaternion_multiply
from embodied_gen.utils.log import logger
from embodied_gen.utils.process_media import alpha_blend_rgba
from embodied_gen.utils.simulation import (
    SIM_COORD_ALIGN,
    load_assets_from_layout_file,
)

__all__ = ["PickEmbodiedGen"]


@register_env("PickEmbodiedGen-v1", max_episode_steps=100)
class PickEmbodiedGen(BaseEnv):
    """PickEmbodiedGen as gym env example for object pick-and-place tasks.

    This environment simulates a robot interacting with 3D assets in the
    embodiedgen generated scene in SAPIEN. It supports multi-environment setups,
    dynamic reconfiguration, and hybrid rendering with 3D Gaussian Splatting.

    Example:
        Use `gym.make` to create the `PickEmbodiedGen-v1` parallel environment.
        ```python
        import gymnasium as gym
        env = gym.make(
            "PickEmbodiedGen-v1",
            num_envs=cfg.num_envs,
            render_mode=cfg.render_mode,
            enable_shadow=cfg.enable_shadow,
            layout_file=cfg.layout_file,
            control_mode=cfg.control_mode,
            camera_cfg=dict(
                camera_eye=cfg.camera_eye,
                camera_target_pt=cfg.camera_target_pt,
                image_hw=cfg.image_hw,
                fovy_deg=cfg.fovy_deg,
            ),
        )
        ```
    """

    SUPPORTED_ROBOTS = ["panda", "panda_wristcam", "fetch"]
    goal_thresh = 0.0

    def __init__(
        self,
        *args,
        robot_uids: str | list[str] = "panda",
        robot_init_qpos_noise: float = 0.02,
        num_envs: int = 1,
        reconfiguration_freq: int = None,
        **kwargs,
    ):
        """Initializes the PickEmbodiedGen environment.

        Args:
            *args: Variable length argument list for the base class.
            robot_uids: The robot(s) to use in the environment.
            robot_init_qpos_noise: Noise added to the robot's initial joint
                positions.
            num_envs: The number of parallel environments to create.
            reconfiguration_freq: How often to reconfigure the scene. If None,
                it is set based on num_envs.
            **kwargs: Additional keyword arguments for environment setup,
                including layout_file, replace_objs, enable_grasp, etc.
        """
        self.robot_init_qpos_noise = robot_init_qpos_noise
        if reconfiguration_freq is None:
            if num_envs == 1:
                reconfiguration_freq = 1
            else:
                reconfiguration_freq = 0

        # Init params from kwargs.
        layout_file = kwargs.pop("layout_file", None)
        replace_objs = kwargs.pop("replace_objs", True)
        self.enable_grasp = kwargs.pop("enable_grasp", False)
        self.init_3dgs_quat = kwargs.pop(
            "init_3dgs_quat", [0.7071, 0, 0, 0.7071]
        )
        # Add small offset in z-axis to avoid collision.
        self.objs_z_offset = kwargs.pop("objs_z_offset", 0.002)
        self.robot_z_offset = kwargs.pop("robot_z_offset", 0.002)
        self.camera_cfg = kwargs.pop("camera_cfg", None)
        if self.camera_cfg is None:
            self.camera_cfg = dict(
                camera_eye=[0.9, 0.0, 1.1],
                camera_target_pt=[0.0, 0.0, 0.9],
                image_hw=[256, 256],
                fovy_deg=75,
            )

        self.layouts = self.init_env_layouts(
            layout_file, num_envs, replace_objs
        )
        self.robot_pose = self.compute_robot_init_pose(
            self.layouts, num_envs, self.robot_z_offset
        )
        self.env_actors = dict()
        self.image_transform = transforms.PILToTensor()

        super().__init__(
            *args,
            robot_uids=robot_uids,
            reconfiguration_freq=reconfiguration_freq,
            num_envs=num_envs,
            **kwargs,
        )

        self.bg_images = dict()
        if self.render_mode == "hybrid":
            self.bg_images = self.render_gs3d_images(
                self.layouts, num_envs, self.init_3dgs_quat
            )

    @staticmethod
    def init_env_layouts(
        layout_file: str, num_envs: int, replace_objs: bool
    ) -> list[LayoutInfo]:
        """Initializes and saves layout files for each environment instance.

        For each environment, this method creates a layout configuration. If
        `replace_objs` is True, it generates new object placements for each
        subsequent environment. The generated layouts are saved as new JSON
        files.

        Args:
            layout_file: Path to the base layout JSON file.
            num_envs: The number of environments to create layouts for.
            replace_objs: If True, generates new object placements for each
                environment after the first one using BFS placement.

        Returns:
            A list of file paths to the generated layout for each environment.
        """
        layouts = []
        for env_idx in range(num_envs):
            if replace_objs and env_idx > 0:
                layout_info = bfs_placement(layout_file)
            else:
                layout_info = json.load(open(layout_file, "r"))
                layout_info = LayoutInfo.from_dict(layout_info)

            layout_path = layout_file.replace(".json", f"_env{env_idx}.json")
            with open(layout_path, "w") as f:
                json.dump(layout_info.to_dict(), f, indent=4)

            layouts.append(layout_path)

        return layouts

    @staticmethod
    def compute_robot_init_pose(
        layouts: list[str], num_envs: int, z_offset: float = 0.0
    ) -> list[list[float]]:
        """Computes the initial pose for the robot in each environment.

        Args:
            layouts: A list of file paths to the environment layouts.
            num_envs: The number of environments.
            z_offset: An optional vertical offset to apply to the robot's
                position to prevent collisions.

        Returns:
            A list of initial poses ([x, y, z, qw, qx, qy, qz]) for the robot
            in each environment.
        """
        robot_pose = []
        for env_idx in range(num_envs):
            layout = json.load(open(layouts[env_idx], "r"))
            layout = LayoutInfo.from_dict(layout)
            robot_node = layout.relation[Scene3DItemEnum.ROBOT.value]
            x, y, z, qx, qy, qz, qw = layout.position[robot_node]
            robot_pose.append([x, y, z + z_offset, qw, qx, qy, qz])

        return robot_pose

    @property
    def _default_sim_config(self):
        """Returns the default simulation configuration.

        Returns:
            The default simulation configuration object.
        """
        return SimConfig(
            scene_config=SceneConfig(
                solver_position_iterations=30,
                # contact_offset=0.04,
                # rest_offset=0.001,
            ),
            # sim_freq=200,
            control_freq=50,
            gpu_memory_config=GPUMemoryConfig(
                max_rigid_contact_count=2**20, max_rigid_patch_count=2**19
            ),
        )

    @property
    def _default_sensor_configs(self):
        """Returns the default sensor configurations for the agent.

        Returns:
            A list containing the default camera configuration.
        """
        pose = sapien_utils.look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1])

        return [
            CameraConfig("base_camera", pose, 128, 128, np.pi / 2, 0.01, 100)
        ]

    @property
    def _default_human_render_camera_configs(self):
        """Returns the default camera configuration for human-friendly rendering.

        Returns:
            The default camera configuration for the renderer.
        """
        pose = sapien_utils.look_at(
            eye=self.camera_cfg["camera_eye"],
            target=self.camera_cfg["camera_target_pt"],
        )

        return CameraConfig(
            "render_camera",
            pose,
            self.camera_cfg["image_hw"][1],
            self.camera_cfg["image_hw"][0],
            np.deg2rad(self.camera_cfg["fovy_deg"]),
            0.01,
            100,
        )

    def _load_agent(self, options: dict):
        """Loads the agent (robot) and a ground plane into the scene.

        Args:
            options: A dictionary of options for loading the agent.
        """
        self.ground = build_ground(self.scene)
        super()._load_agent(options, sapien.Pose(p=[-10, 0, 10]))

    def _load_scene(self, options: dict):
        """Loads all assets, objects, and the goal site into the scene.

        This method iterates through the layouts for each environment, loads the
        specified assets, and adds them to the simulation. It also creates a
        kinematic sphere to represent the goal site.

        Args:
            options: A dictionary of options for loading the scene.
        """
        all_objects = []
        logger.info(f"Loading EmbodiedGen assets...")
        for env_idx in range(self.num_envs):
            env_actors = load_assets_from_layout_file(
                self.scene,
                self.layouts[env_idx],
                z_offset=self.objs_z_offset,
                env_idx=env_idx,
            )
            self.env_actors[f"env{env_idx}"] = env_actors
            all_objects.extend(env_actors.values())

        self.obj = all_objects[-1]
        for obj in all_objects:
            self.remove_from_state_dict_registry(obj)

        self.all_objects = Actor.merge(all_objects, name="all_objects")
        self.add_to_state_dict_registry(self.all_objects)

        self.goal_site = actors.build_sphere(
            self.scene,
            radius=self.goal_thresh,
            color=[0, 1, 0, 0],
            name="goal_site",
            body_type="kinematic",
            add_collision=False,
            initial_pose=sapien.Pose(),
        )
        self._hidden_objects.append(self.goal_site)

    def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
        """Initializes an episode for a given set of environments.

        This method sets the goal position, resets the robot's joint positions
        with optional noise, and sets its root pose.

        Args:
            env_idx: A tensor of environment indices to initialize.
            options: A dictionary of options for initialization.
        """
        with torch.device(self.device):
            b = len(env_idx)
            goal_xyz = torch.zeros((b, 3))
            goal_xyz[:, :2] = torch.rand((b, 2)) * 0.2 - 0.1
            self.goal_site.set_pose(Pose.create_from_pq(goal_xyz))

            qpos = np.array(
                [
                    0.0,
                    np.pi / 8,
                    0,
                    -np.pi * 3 / 8,
                    0,
                    np.pi * 3 / 4,
                    np.pi / 4,
                    0.04,
                    0.04,
                ]
            )
            qpos = (
                np.random.normal(
                    0, self.robot_init_qpos_noise, (self.num_envs, len(qpos))
                )
                + qpos
            )
            qpos[:, -2:] = 0.04
            self.agent.robot.set_root_pose(np.array(self.robot_pose))
            self.agent.reset(qpos)
            self.agent.init_qpos = qpos
            self.agent.controller.controllers["gripper"].reset()

    def render_gs3d_images(
        self, layouts: list[str], num_envs: int, init_quat: list[float]
    ) -> dict[str, np.ndarray]:
        """Renders background images using a pre-trained Gaussian Splatting model.

        This method pre-renders the static background for each environment from
        the perspective of all cameras to be used for hybrid rendering.

        Args:
            layouts: A list of file paths to the environment layouts.
            num_envs: The number of environments.
            init_quat: An initial quaternion to orient the Gaussian Splatting
                model.

        Returns:
            A dictionary mapping a unique key (e.g., 'camera-env_idx') to the
            rendered background image as a numpy array.
        """
        sim_coord_align = (
            torch.tensor(SIM_COORD_ALIGN).to(torch.float32).to(self.device)
        )
        cameras = self.scene.sensors.copy()
        cameras.update(self.scene.human_render_cameras)

        # Preload the background Gaussian Splatting model.
        asset_root = os.path.dirname(layouts[0])
        layout = LayoutInfo.from_dict(json.load(open(layouts[0], "r")))
        bg_node = layout.relation[Scene3DItemEnum.BACKGROUND.value]
        gs_path = os.path.join(
            asset_root, layout.assets[bg_node], "gs_model.ply"
        )
        raw_gs: GaussianOperator = GaussianOperator.load_from_ply(gs_path)
        bg_images = dict()
        for env_idx in tqdm(range(num_envs), desc="Pre-rendering Background"):
            layout = json.load(open(layouts[env_idx], "r"))
            layout = LayoutInfo.from_dict(layout)
            x, y, z, qx, qy, qz, qw = layout.position[bg_node]
            qx, qy, qz, qw = quaternion_multiply([qx, qy, qz, qw], init_quat)
            init_pose = torch.tensor([x, y, z, qx, qy, qz, qw])
            gs_model = raw_gs.get_gaussians(instance_pose=init_pose)
            for key in cameras:
                camera = cameras[key]
                Ks = camera.camera.get_intrinsic_matrix()  # (n_env, 3, 3)
                c2w = camera.camera.get_model_matrix()  # (n_env, 4, 4)
                result = gs_model.render(
                    c2w[env_idx] @ sim_coord_align,
                    Ks[env_idx],
                    image_width=camera.config.width,
                    image_height=camera.config.height,
                )
                bg_images[f"{key}-env{env_idx}"] = result.rgb[..., ::-1]

        return bg_images

    def render(self):
        """Renders the environment based on the configured render_mode.

        Raises:
            RuntimeError: If `render_mode` is not set.
            NotImplementedError: If the `render_mode` is not supported.

        Returns:
            The rendered output, which varies depending on the render mode.
        """
        if self.render_mode is None:
            raise RuntimeError("render_mode is not set.")
        if self.render_mode == "human":
            return self.render_human()
        elif self.render_mode == "rgb_array":
            res = self.render_rgb_array()
            return res
        elif self.render_mode == "sensors":
            res = self.render_sensors()
            return res
        elif self.render_mode == "all":
            return self.render_all()
        elif self.render_mode == "hybrid":
            return self.hybrid_render()
        else:
            raise NotImplementedError(
                f"Unsupported render mode {self.render_mode}."
            )

    def render_rgb_array(
        self, camera_name: str = None, return_alpha: bool = False
    ):
        """Renders an RGB image from the human-facing render camera.

        Args:
            camera_name: The name of the camera to render from. If None, uses
                all human render cameras.
            return_alpha: Whether to include the alpha channel in the output.

        Returns:
            A numpy array representing the rendered image(s). If multiple
            cameras are used, the images are tiled.
        """
        for obj in self._hidden_objects:
            obj.show_visual()
        self.scene.update_render(
            update_sensors=False, update_human_render_cameras=True
        )
        images = []
        render_images = self.scene.get_human_render_camera_images(
            camera_name, return_alpha
        )
        for image in render_images.values():
            images.append(image)
        if len(images) == 0:
            return None
        if len(images) == 1:
            return images[0]
        for obj in self._hidden_objects:
            obj.hide_visual()
        return tile_images(images)

    def render_sensors(self):
        """Renders images from all on-board sensor cameras.

        Returns:
            A tiled image of all sensor outputs as a numpy array.
        """
        images = []
        sensor_images = self.get_sensor_images()
        for image in sensor_images.values():
            for img in image.values():
                images.append(img)
        return tile_images(images)

    def hybrid_render(self):
        """Renders a hybrid image by blending simulated foreground with a background.

        The foreground is rendered with an alpha channel and then blended with
        the pre-rendered Gaussian Splatting background image.

        Returns:
            A torch tensor of the final blended RGB images.
        """
        fg_images = self.render_rgb_array(
            return_alpha=True
        )  # (n_env, h, w, 3)
        images = []
        for key in self.bg_images:
            if "render_camera" not in key:
                continue
            env_idx = int(key.split("-env")[-1])
            rgba = alpha_blend_rgba(
                fg_images[env_idx].cpu().numpy(), self.bg_images[key]
            )
            images.append(self.image_transform(rgba))

        images = torch.stack(images, dim=0)
        images = images.permute(0, 2, 3, 1)

        return images[..., :3]

    def evaluate(self):
        """Evaluates the current state of the environment.

        Checks for task success criteria such as whether the object is grasped,
        placed at the goal, and if the robot is static.

        Returns:
            A dictionary containing boolean tensors for various success
            metrics, including 'is_grasped', 'is_obj_placed', and overall
            'success'.
        """
        obj_to_goal_pos = (
            self.obj.pose.p
        )  # self.goal_site.pose.p - self.obj.pose.p
        is_obj_placed = (
            torch.linalg.norm(obj_to_goal_pos, axis=1) <= self.goal_thresh
        )
        is_grasped = self.agent.is_grasping(self.obj)
        is_robot_static = self.agent.is_static(0.2)

        return dict(
            is_grasped=is_grasped,
            obj_to_goal_pos=obj_to_goal_pos,
            is_obj_placed=is_obj_placed,
            is_robot_static=is_robot_static,
            is_grasping=self.agent.is_grasping(self.obj),
            success=torch.logical_and(is_obj_placed, is_robot_static),
        )

    def _get_obs_extra(self, info: dict):
        """Gets extra information for the observation dictionary.

        Args:
            info: A dictionary containing evaluation information.

        Returns:
            An empty dictionary, as no extra observations are added.
        """

        return dict()

    def compute_dense_reward(self, obs: any, action: torch.Tensor, info: dict):
        """Computes a dense reward for the current step.

        The reward is a composite of reaching, grasping, placing, and
        maintaining a static final pose.

        Args:
            obs: The current observation.
            action: The action taken in the current step.
            info: A dictionary containing evaluation information from `evaluate()`.

        Returns:
            A tensor containing the dense reward for each environment.
        """
        tcp_to_obj_dist = torch.linalg.norm(
            self.obj.pose.p - self.agent.tcp.pose.p, axis=1
        )
        reaching_reward = 1 - torch.tanh(5 * tcp_to_obj_dist)
        reward = reaching_reward

        is_grasped = info["is_grasped"]
        reward += is_grasped

        # obj_to_goal_dist = torch.linalg.norm(
        #     self.goal_site.pose.p - self.obj.pose.p, axis=1
        # )
        obj_to_goal_dist = torch.linalg.norm(
            self.obj.pose.p - self.obj.pose.p, axis=1
        )
        place_reward = 1 - torch.tanh(5 * obj_to_goal_dist)
        reward += place_reward * is_grasped

        reward += info["is_obj_placed"] * is_grasped

        static_reward = 1 - torch.tanh(
            5
            * torch.linalg.norm(self.agent.robot.get_qvel()[..., :-2], axis=1)
        )
        reward += static_reward * info["is_obj_placed"] * is_grasped

        reward[info["success"]] = 6
        return reward

    def compute_normalized_dense_reward(
        self, obs: any, action: torch.Tensor, info: dict
    ):
        """Computes a dense reward normalized to be between 0 and 1.

        Args:
            obs: The current observation.
            action: The action taken in the current step.
            info: A dictionary containing evaluation information from `evaluate()`.

        Returns:
            A tensor containing the normalized dense reward for each environment.
        """
        return self.compute_dense_reward(obs=obs, action=action, info=info) / 6