diff --git a/mani_skill2/envs/__init__.py b/mani_skill2/envs/__init__.py index 10d14cba1..39c16f81e 100644 --- a/mani_skill2/envs/__init__.py +++ b/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 * diff --git a/mani_skill2/envs/assembly/assembling_kits.py b/mani_skill2/envs/assembly/assembling_kits.py index 0b7d96bae..17a434e2f 100644 --- a/mani_skill2/envs/assembly/assembling_kits.py +++ b/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) diff --git a/mani_skill2/envs/assembly/base_env.py b/mani_skill2/envs/assembly/base_env.py index b7aa6047e..8be83eb19 100644 --- a/mani_skill2/envs/assembly/base_env.py +++ b/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( diff --git a/mani_skill2/envs/assembly/peg_insertion_side.py b/mani_skill2/envs/assembly/peg_insertion_side.py index 6efa4e603..28818453b 100644 --- a/mani_skill2/envs/assembly/peg_insertion_side.py +++ b/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 diff --git a/mani_skill2/envs/assembly/plug_charger.py b/mani_skill2/envs/assembly/plug_charger.py index d8a2f61aa..3de4e7d98 100644 --- a/mani_skill2/envs/assembly/plug_charger.py +++ b/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, diff --git a/mani_skill2/envs/misc/assets/Sink_19.glb b/mani_skill2/envs/misc/assets/Sink_19.glb new file mode 100644 index 000000000..8035f974f Binary files /dev/null and b/mani_skill2/envs/misc/assets/Sink_19.glb differ diff --git a/mani_skill2/envs/misc/turn_faucet.py b/mani_skill2/envs/misc/turn_faucet.py index 9d04dbe9b..a60ec829f 100644 --- a/mani_skill2/envs/misc/turn_faucet.py +++ b/mani_skill2/envs/misc/turn_faucet.py @@ -1,3 +1,4 @@ +import os.path as osp from collections import OrderedDict from pathlib import Path from typing import Dict, List, Union @@ -5,6 +6,7 @@ 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, diff --git a/mani_skill2/envs/ms1/move_bucket.py b/mani_skill2/envs/ms1/move_bucket.py index 8904a3615..b00bbae7e 100644 --- a/mani_skill2/envs/ms1/move_bucket.py +++ b/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 diff --git a/mani_skill2/envs/ms1/open_cabinet_door_drawer.py b/mani_skill2/envs/ms1/open_cabinet_door_drawer.py index 2f09b766a..463a4707e 100644 --- a/mani_skill2/envs/ms1/open_cabinet_door_drawer.py +++ b/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 diff --git a/mani_skill2/envs/ms1/push_chair.py b/mani_skill2/envs/ms1/push_chair.py index 584233422..26ef7c76a 100644 --- a/mani_skill2/envs/ms1/push_chair.py +++ b/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 diff --git a/mani_skill2/envs/pick_and_place/base_env.py b/mani_skill2/envs/pick_and_place/base_env.py index 264bae98d..a1c083e6d 100644 --- a/mani_skill2/envs/pick_and_place/base_env.py +++ b/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() diff --git a/mani_skill2/envs/pick_and_place/pick_clutter.py b/mani_skill2/envs/pick_and_place/pick_clutter.py index 1305fcc53..23191a0ef 100644 --- a/mani_skill2/envs/pick_and_place/pick_clutter.py +++ b/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 = [] diff --git a/mani_skill2/envs/pick_and_place/pick_cube.py b/mani_skill2/envs/pick_and_place/pick_cube.py index 9d41519a8..79a1143b9 100644 --- a/mani_skill2/envs/pick_and_place/pick_cube.py +++ b/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) diff --git a/mani_skill2/envs/pick_and_place/pick_single.py b/mani_skill2/envs/pick_and_place/pick_single.py index 37beec512..dcc839067 100644 --- a/mani_skill2/envs/pick_and_place/pick_single.py +++ b/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) diff --git a/mani_skill2/envs/pick_and_place/stack_cube.py b/mani_skill2/envs/pick_and_place/stack_cube.py index 2094273f6..11e249397 100644 --- a/mani_skill2/envs/pick_and_place/stack_cube.py +++ b/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") diff --git a/mani_skill2/envs/sapien_env.py b/mani_skill2/envs/sapien_env.py index f06f1f3ba..54c18ab87 100644 --- a/mani_skill2/envs/sapien_env.py +++ b/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 diff --git a/mani_skill2/envs/scenes/base_env.py b/mani_skill2/envs/scenes/base_env.py new file mode 100644 index 000000000..8026170c8 --- /dev/null +++ b/mani_skill2/envs/scenes/base_env.py @@ -0,0 +1,111 @@ +import numpy as np +import sapien as sapien +import sapien.physx as physx +from sapien import Pose + +from mani_skill2.agents.robots import Panda +from mani_skill2.envs.sapien_env import BaseEnv +from mani_skill2.sensors.camera import CameraConfig +from mani_skill2.utils.sapien_utils import look_at, vectorize_pose +from mani_skill2.utils.scene_builder import SceneBuilder +from mani_skill2.utils.scene_builder.ai2thor import ( + ArchitecTHORSceneBuilder, + ProcTHORSceneBuilder, + RoboTHORSceneBuilder, + iTHORSceneBuilder, +) + + +class SceneManipulationEnv(BaseEnv): + agent: Panda + """ + Args: + robot_uid: Which robot to place into the scene. Default is "panda" + + fixed_scene: whether to sample a single scene and never reconfigure the scene during episode resets + Default to True as reconfiguration/reloading scenes is expensive. When true, call env.reset(seed=seed, options=dict(reconfigure=True)) + + scene_builder_cls: Scene builder class to build a scene with. Default is the ArchitecTHORSceneBuilder which builds a scene from AI2THOR. + Any of the AI2THOR SceneBuilders are supported in this environment + + convex_decomposition: Choice of convex decomposition algorithm to generate collision meshes for objects. Default is `coacd` which uses https://github.com/SarahWeiii/CoACD + """ + + def __init__( + self, + *args, + robot_uid="panda", + robot_init_qpos_noise=0.02, + fixed_scene=True, + scene_builder_cls: SceneBuilder = ArchitecTHORSceneBuilder, + convex_decomposition="coacd", + **kwargs + ): + self.robot_init_qpos_noise = robot_init_qpos_noise + self.fixed_scene = fixed_scene + self.sampled_scene_idx: int = None + self.scene_builder = scene_builder_cls() + self.scene_ids = np.arange(0, len(self.scene_builder.scene_configs)) + self.convex_decomposition = convex_decomposition + super().__init__(*args, robot_uid=robot_uid, **kwargs) + + def reset(self, seed=None, options=None): + self._set_episode_rng(seed) + if options is None: + options = dict(reconfigure=True) + if not self.fixed_scene: + options["reconfigure"] = True + if options["reconfigure"]: + self.sampled_scene_idx = self._episode_rng.randint(0, len(self.scene_ids)) + return super().reset(seed, options) + + def _load_actors(self): + self.scene_builder.build(self._scene, scene_id=self.sampled_scene_idx, convex_decomposition=self.convex_decomposition) + + def _initialize_agent(self): + if self.robot_uid == "panda": + # fmt: off + # EE at [0.615, 0, 0.17] + qpos = np.array( + [0.0, np.pi / 8, 0, -np.pi * 5 / 8, 0, np.pi * 3 / 4, np.pi / 4, 0.04, 0.04] + ) + # fmt: on + qpos[:-2] += self._episode_rng.normal( + 0, self.robot_init_qpos_noise, len(qpos) - 2 + ) + self.agent.reset(qpos) + self.agent.robot.set_pose(Pose([-0.615, 0, 0])) + elif self.robot_uid == "xmate3_robotiq": + qpos = np.array( + [0, np.pi / 6, 0, np.pi / 3, 0, np.pi / 2, -np.pi / 2, 0, 0] + ) + qpos[:-2] += self._episode_rng.normal( + 0, self.robot_init_qpos_noise, len(qpos) - 2 + ) + self.agent.reset(qpos) + self.agent.robot.set_pose(Pose([-0.562, 0, 0])) + else: + raise NotImplementedError(self.robot_uid) + + def _register_sensors(self): + pose = look_at([0.3, 0, 0.6], [-0.1, 0, 0.1]) + return CameraConfig( + "base_camera", pose.p, pose.q, 128, 128, np.pi / 2, 0.01, 10 + ) + + 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]) + 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) + + def _get_obs_agent(self): + obs = self.agent.get_proprioception() + obs["base_pose"] = vectorize_pose(self.agent.robot.pose) + return obs diff --git a/mani_skill2/envs/scenes/pick_object.py b/mani_skill2/envs/scenes/pick_object.py new file mode 100644 index 000000000..c85fe0f31 --- /dev/null +++ b/mani_skill2/envs/scenes/pick_object.py @@ -0,0 +1,59 @@ +from collections import OrderedDict + +import numpy as np +import sapien +import sapien.physx as physx + +from mani_skill2.utils.registration import register_env +from mani_skill2.utils.sapien_utils import vectorize_pose + +from .base_env import SceneManipulationEnv + + +@register_env("PickObjectScene-v0", max_episode_steps=200) +class PickObjectSceneEnv(SceneManipulationEnv): + """ + Args: + + """ + + def __init__(self, *args, **kwargs): + self.box_length = 0.1 + super().__init__(*args, **kwargs) + + def _load_actors(self): + super()._load_actors() + + def reconfigure(self): + super().reconfigure() + self.init_state = self.get_state() + + def _get_obs_extra(self) -> OrderedDict: + obs = OrderedDict( + tcp_pose=vectorize_pose(self.agent.tcp.pose), + obj_pose=vectorize_pose(self.goal_obj.pose), + ) + return obs + + def evaluate(self, **kwargs): + return dict(success=self.agent.is_grasping(self.goal_obj)) + + def compute_dense_reward(self, info, **kwargs): + return 1 - np.tanh(np.linalg.norm(self.goal_obj.pose.p - self.agent.tcp.pose.p)) + + def compute_normalized_dense_reward(self, **kwargs): + return self.compute_dense_reward(**kwargs) / 1 + + def _initialize_actors(self): + self.set_state(self.init_state) + + def _initialize_task(self): + # pick a random goal object to pick up + self.goal_obj: sapien.Entity = self._episode_rng.choice( + self.scene_builder.movable_objects + ) + print(f"Target Object: {self.goal_obj.name}") + + def reset(self, seed=None, options=None): + # sample a scene + return super().reset(seed, options)