merge new env code, refactors + new scene builder api

mani_skill2/envs/__init__.py

@@ -1,5 +1,7 @@
 from .assembly import *
+from .composer.tasks import *
 from .misc import *
 from .mpm import *
 from .ms1 import *
 from .pick_and_place import *
+from .scenes.pick_object import *

mani_skill2/envs/assembly/assembling_kits.py

@@ -9,6 +9,7 @@
 from mani_skill2.utils.io_utils import load_json
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import look_at, vectorize_pose
+from mani_skill2.utils.scene_builder import TableSceneBuilder
 
 from .base_env import StationaryManipulationEnv
 
@@ -106,7 +107,7 @@ def _load_object(self, object_id):
         return builder.build(f"obj_{object_id:02d}")
 
     def _load_actors(self):
-        self._add_ground(render=self.bg_name is None)
+        TableSceneBuilder().build(self._scene)
 
         self.kit = self._load_kit()
         self.obj = self._load_object(self.object_id)

mani_skill2/envs/assembly/base_env.py

@@ -27,7 +27,7 @@ def _get_default_scene_config(self):
         return scene_config
 
     def _initialize_agent(self):
-        if self.robot_uid == "panda":
+        if self.robot_uid == "panda" or self.robot_uid == "panda_realsensed435":
             # fmt: off
             # EE at [0.615, 0, 0.17]
             qpos = np.array(

mani_skill2/envs/assembly/peg_insertion_side.py

@@ -7,6 +7,7 @@
 
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import hex2rgba, look_at, vectorize_pose
+from mani_skill2.utils.scene_builder import TableSceneBuilder
 
 from .base_env import StationaryManipulationEnv
 
@@ -56,7 +57,7 @@ def _build_box_with_hole(
         return builder.build_static(name)
 
     def _load_actors(self):
-        self._add_ground(render=self.bg_name is None)
+        TableSceneBuilder().build(self._scene)
 
         # peg
         # length, radius = 0.1, 0.02

mani_skill2/envs/assembly/plug_charger.py

@@ -9,6 +9,7 @@
 
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import hex2rgba, look_at, vectorize_pose
+from mani_skill2.utils.scene_builder import TableSceneBuilder
 
 from .base_env import StationaryManipulationEnv
 
@@ -96,7 +97,7 @@ def _build_receptacle(self, peg_size, receptacle_size, gap):
         return builder.build_static(name="receptacle")
 
     def _load_actors(self):
-        self._add_ground(render=self.bg_name is None)
+        TableSceneBuilder().build(self._scene)
         self.charger = self._build_charger(
             self._peg_size,
             self._base_size,

mani_skill2/envs/misc/turn_faucet.py

@@ -1,10 +1,12 @@
+import os.path as osp
 from collections import OrderedDict
 from pathlib import Path
 from typing import Dict, List, Union
 
 import numpy as np
 import sapien
 import sapien.physx as physx
+import sapien.render
 import trimesh
 import trimesh.sample
 from sapien import Pose
@@ -15,8 +17,10 @@
 from mani_skill2.agents.robots.panda.panda import Panda
 from mani_skill2.envs.sapien_env import BaseEnv
 from mani_skill2.sensors.camera import CameraConfig
+from mani_skill2.utils.building.ground import build_tesselated_square_floor
 from mani_skill2.utils.common import np_random, random_choice
 from mani_skill2.utils.geometry import transform_points
+from mani_skill2.utils.geometry.trimesh_utils import get_component_mesh
 from mani_skill2.utils.io_utils import load_json
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import (
@@ -26,7 +30,6 @@
     set_articulation_render_material,
     vectorize_pose,
 )
-from mani_skill2.utils.trimesh_utils import get_component_mesh
 
 
 class TurnFaucetBaseEnv(BaseEnv):
@@ -43,7 +46,43 @@ def __init__(
         super().__init__(*args, robot_uid=robot_uid, **kwargs)
 
     def _load_actors(self):
-        self._add_ground(render=self.bg_name is None)
+        # builder = self._scene.create_actor_builder()
+        # model_dir = Path(osp.dirname(__file__)) / "assets"
+        # scale = 1
+        # collision_file = str(model_dir / "Sink_19.glb")  # a metal table
+        # sink_pose = sapien.Pose(q=euler2quat(np.pi / 2, 0, 0))
+        # builder.add_nonconvex_collision_from_file(
+        #     filename=collision_file, scale=[scale] * 3, material=None, pose=sink_pose
+        # )
+        # visual_file = str(model_dir / "Sink_19.glb")
+        # builder.add_visual_from_file(
+        #     filename=visual_file, scale=[scale] * 3, pose=sink_pose
+        # )
+        # self.sink = builder.build_static(name="sink")
+        # aabb = self.sink.find_component_by_type(
+        #     sapien.render.RenderBodyComponent
+        # ).compute_global_aabb_tight()
+        # sink_height = aabb[1, 2] - aabb[0, 2]
+
+        # self.sink.set_pose(
+        #     Pose(p=[-0.24, 0, -sink_height], q=euler2quat(0, 0, -np.pi / 2))
+        # )
+
+        build_tesselated_square_floor(self._scene)
+
+        # # add wall
+        # wall_mtl = sapien.render.RenderMaterial(
+        #     base_color=[32 / 255, 67 / 255, 80 / 255, 1],
+        #     metallic=0,
+        #     roughness=0.9,
+        #     specular=0.8,
+        # )
+        # wall = self._scene.create_actor_builder()
+        # half_size = (0.02, 6, 2.1)
+        # wall.add_box_collision(half_size=half_size)
+        # wall.add_box_visual(half_size=half_size, material=wall_mtl)
+        # self.wall = wall.build_static("wall")
+        # self.wall.set_pose(Pose(p=[0.25, 0, 1]))
 
     def _initialize_agent(self):
         if self.robot_uid == "panda":
@@ -70,13 +109,12 @@ def _register_sensors(self):
         )
 
     def _register_render_cameras(self):
-        pose = look_at([0.5, 0.5, 1.0], [0.0, 0.0, 0.5])
+        pose = look_at([-1.3, 0.6, 0.6], [0.0, 0.0, 0.4])
         return CameraConfig("render_camera", pose.p, pose.q, 512, 512, 1, 0.01, 10)
 
     def _setup_viewer(self):
         super()._setup_viewer()
-        self._viewer.set_camera_xyz(1.0, 0.0, 1.2)
-        self._viewer.set_camera_rpy(0, -0.5, 3.14)
+        self._viewer.set_camera_pose(look_at([-1.3, 0.6, 0.6], [0.0, 0.0, 0.4]))
 
 
 @register_env("TurnFaucet-v0", max_episode_steps=200)
@@ -199,6 +237,7 @@ def _load_faucet(self):
         loader.set_density(density)
         articulation.set_name("faucet")
 
+        # TODO (stao): find out why we did this before
         set_articulation_render_material(
             articulation, color=hex2rgba("#AAAAAA"), metallic=1, roughness=0.4
         )
@@ -250,6 +289,8 @@ def _initialize_articulations(self):
         ori = self._episode_rng.uniform(-np.pi / 12, np.pi / 12)
         q = euler2quat(0, 0, ori)
         self.faucet.set_pose(Pose(p, q))
+        # self.sink.set_pose(self.sink.pose * Pose(p=[p[0], p[1], 0.011], q=q))
+        # self.wall.set_pose(self.wall.pose * Pose(p=[p[0], p[1], 0], q=q))
 
     def _initialize_task(self):
         self._set_target_link()
@@ -319,7 +360,7 @@ def _set_init_and_target_angle(self):
 
     def _get_obs_extra(self) -> OrderedDict:
         obs = OrderedDict(
-            tcp_pose=vectorize_pose(self.tcp.pose),
+            tcp_pose=vectorize_pose(self.agent.tcp.pose),
             target_angle_diff=np.array(self.target_angle_diff),
             target_joint_axis=self.target_joint_axis,
             target_link_pos=self.target_link_pos,

mani_skill2/envs/ms1/move_bucket.py

@@ -16,9 +16,9 @@
     get_local_axis_aligned_bbox_for_link,
     transform_points,
 )
+from mani_skill2.utils.geometry.trimesh_utils import get_actor_visual_mesh
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import get_obj_by_name, vectorize_pose
-from mani_skill2.utils.trimesh_utils import get_actor_visual_mesh
 
 from .base_env import MS1BaseEnv
 

mani_skill2/envs/ms1/open_cabinet_door_drawer.py

@@ -12,14 +12,14 @@
 from mani_skill2.agents.robots.mobile_panda import MobilePandaSingleArm
 from mani_skill2.utils.common import np_random, random_choice
 from mani_skill2.utils.geometry import angle_distance, transform_points
-from mani_skill2.utils.registration import register_env
-from mani_skill2.utils.sapien_utils import get_obj_by_name, vectorize_pose
-from mani_skill2.utils.trimesh_utils import (
+from mani_skill2.utils.geometry.trimesh_utils import (
     get_articulation_meshes,
     get_render_body_meshes,
     get_render_shape_meshes,
     merge_meshes,
 )
+from mani_skill2.utils.registration import register_env
+from mani_skill2.utils.sapien_utils import get_obj_by_name, vectorize_pose
 
 from .base_env import MS1BaseEnv
 

mani_skill2/envs/ms1/push_chair.py

@@ -10,9 +10,9 @@
 from mani_skill2.agents.robots.mobile_panda import MobilePandaDualArm
 from mani_skill2.utils.common import np_random
 from mani_skill2.utils.geometry import transform_points
+from mani_skill2.utils.geometry.trimesh_utils import get_actor_visual_mesh
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import get_obj_by_name, vectorize_pose
-from mani_skill2.utils.trimesh_utils import get_actor_visual_mesh
 
 from .base_env import MS1BaseEnv
 

mani_skill2/envs/pick_and_place/base_env.py

@@ -106,16 +106,13 @@ def _register_sensors(self):
         )
 
     def _register_render_cameras(self):
-        if self.robot_uid == "panda":
-            pose = look_at([0.4, 0.4, 0.8], [0.0, 0.0, 0.4])
-        else:
-            pose = look_at([0.5, 0.5, 1.0], [0.0, 0.0, 0.5])
+        pose = look_at([0.6, 0.7, 0.6], [0.0, 0.0, 0.35])
         return CameraConfig("render_camera", pose.p, pose.q, 512, 512, 1, 0.01, 10)
 
     def _setup_viewer(self):
         super()._setup_viewer()
-        self._viewer.set_camera_xyz(0.8, 0, 1.0)
-        self._viewer.set_camera_rpy(0, -0.5, 3.14)
+        pose = look_at([0.6, 0.7, 0.6], [0.0, 0.0, 0.35])
+        self._viewer.set_camera_pose(pose)
 
     def _get_obs_agent(self):
         obs = self.agent.get_proprioception()

mani_skill2/envs/pick_and_place/pick_clutter.py

@@ -16,6 +16,7 @@
     set_entity_visibility,
     vectorize_pose,
 )
+from mani_skill2.utils.scene_builder import TableSceneBuilder
 
 from .base_env import StationaryManipulationEnv
 from .pick_single import PickSingleYCBEnv, build_actor_ycb
@@ -65,7 +66,7 @@ def __init__(
         super().__init__(**kwargs)
 
     def _load_actors(self):
-        self._add_ground(render=self.bg_name is None)
+        TableSceneBuilder().build(self._scene)
 
         self.objs: List[sapien.Entity] = []
         self.bbox_sizes = []

mani_skill2/envs/pick_and_place/pick_cube.py

@@ -7,6 +7,7 @@
 
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import hide_entity, show_entity, vectorize_pose
+from mani_skill2.utils.scene_builder import TableSceneBuilder
 
 from .base_env import StationaryManipulationEnv
 
@@ -22,7 +23,7 @@ def __init__(self, *args, obj_init_rot_z=True, **kwargs):
         super().__init__(*args, **kwargs)
 
     def _load_actors(self):
-        self._add_ground(render=self.bg_name is None)
+        TableSceneBuilder().build(self._scene)
         self.obj = self._build_cube(self.cube_half_size)
         self.goal_site = self._build_sphere_site(self.goal_thresh)
 

mani_skill2/envs/pick_and_place/pick_single.py

@@ -20,6 +20,7 @@
     show_entity,
     vectorize_pose,
 )
+from mani_skill2.utils.scene_builder import TableSceneBuilder
 
 from .base_env import StationaryManipulationEnv
 
@@ -80,7 +81,7 @@ def _check_assets(self):
         pass
 
     def _load_actors(self):
-        self._add_ground(render=self.bg_name is None)
+        TableSceneBuilder().build(self._scene)
         self._load_model()
         obj_comp = self.obj.find_component_by_type(physx.PhysxRigidDynamicComponent)
         obj_comp.set_linear_damping(0.1)

mani_skill2/envs/pick_and_place/stack_cube.py

@@ -8,6 +8,7 @@
 
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import check_actor_static, vectorize_pose
+from mani_skill2.utils.scene_builder import TableSceneBuilder
 
 from .base_env import StationaryManipulationEnv
 
@@ -68,7 +69,7 @@ def _get_default_scene_config(self):
         return scene_config
 
     def _load_actors(self):
-        self._add_ground(render=self.bg_name is None)
+        TableSceneBuilder().build(self._scene)
 
         self.box_half_size = np.float32([0.02] * 3)
         self.cubeA = self._build_cube(self.box_half_size, color=(1, 0, 0), name="cubeA")

mani_skill2/envs/sapien_env.py

@@ -22,6 +22,11 @@
 )
 from mani_skill2.sensors.depth_camera import StereoDepthCamera, StereoDepthCameraConfig
 from mani_skill2.utils.common import convert_observation_to_space, flatten_state_dict
+from mani_skill2.utils.geometry.trimesh_utils import (
+    get_articulation_meshes,
+    get_component_meshes,
+    merge_meshes,
+)
 from mani_skill2.utils.sapien_utils import (
     get_actor_state,
     get_articulation_state,
@@ -30,11 +35,6 @@
     set_articulation_render_material,
     set_articulation_state,
 )
-from mani_skill2.utils.trimesh_utils import (
-    get_articulation_meshes,
-    get_component_meshes,
-    merge_meshes,
-)
 from mani_skill2.utils.visualization.misc import observations_to_images, tile_images