from typing import Any, Union
import numpy as np
import sapien
import torch
import mani_skill.envs.utils.randomization as randomization
from mani_skill.agents.robots import SO100, Fetch, Panda, WidowXAI, XArm6Robotiq
from mani_skill.envs.sapien_env import BaseEnv
from mani_skill.envs.tasks.tabletop.pick_cube_cfgs import PICK_CUBE_CONFIGS
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.registration import register_env
from mani_skill.utils.scene_builder.table import TableSceneBuilder
from mani_skill.utils.structs.pose import Pose
[docs]PICK_CUBE_DOC_STRING = """**Task Description:**
A simple task where the objective is to grasp a red cube with the {robot_id} robot and move it to a target goal position. This is also the *baseline* task to test whether a robot with manipulation
capabilities can be simulated and trained properly. Hence there is extra code for some robots to set them up properly in this environment as well as the table scene builder.
**Randomizations:**
- the cube's xy position is randomized on top of a table in the region [0.1, 0.1] x [-0.1, -0.1]. It is placed flat on the table
- the cube's z-axis rotation is randomized to a random angle
- the target goal position (marked by a green sphere) of the cube has its xy position randomized in the region [0.1, 0.1] x [-0.1, -0.1] and z randomized in [0, 0.3]
**Success Conditions:**
- the cube position is within `goal_thresh` (default 0.025m) euclidean distance of the goal position
- the robot is static (q velocity < 0.2)
"""
@register_env("PickCube-v1", max_episode_steps=50)
[docs]class PickCubeEnv(BaseEnv):
[docs] _sample_video_link = "https://github.com/haosulab/ManiSkill/raw/main/figures/environment_demos/PickCube-v1_rt.mp4"
[docs] SUPPORTED_ROBOTS = [
"panda",
"fetch",
"xarm6_robotiq",
"so100",
"widowxai",
]
[docs] agent: Union[Panda, Fetch, XArm6Robotiq, SO100, WidowXAI]
[docs] cube_spawn_half_size = 0.05
[docs] cube_spawn_center = (0, 0)
def __init__(self, *args, robot_uids="panda", robot_init_qpos_noise=0.02, **kwargs):
[docs] self.robot_init_qpos_noise = robot_init_qpos_noise
if robot_uids in PICK_CUBE_CONFIGS:
cfg = PICK_CUBE_CONFIGS[robot_uids]
else:
cfg = PICK_CUBE_CONFIGS["panda"]
[docs] self.cube_half_size = cfg["cube_half_size"]
self.goal_thresh = cfg["goal_thresh"]
self.cube_spawn_half_size = cfg["cube_spawn_half_size"]
self.cube_spawn_center = cfg["cube_spawn_center"]
[docs] self.max_goal_height = cfg["max_goal_height"]
[docs] self.sensor_cam_eye_pos = cfg["sensor_cam_eye_pos"]
[docs] self.sensor_cam_target_pos = cfg["sensor_cam_target_pos"]
[docs] self.human_cam_eye_pos = cfg["human_cam_eye_pos"]
[docs] self.human_cam_target_pos = cfg["human_cam_target_pos"]
super().__init__(*args, robot_uids=robot_uids, **kwargs)
@property
[docs] def _default_sensor_configs(self):
pose = sapien_utils.look_at(
eye=self.sensor_cam_eye_pos, target=self.sensor_cam_target_pos
)
return [CameraConfig("base_camera", pose, 128, 128, np.pi / 2, 0.01, 100)]
@property
[docs] def _default_human_render_camera_configs(self):
pose = sapien_utils.look_at(
eye=self.human_cam_eye_pos, target=self.human_cam_target_pos
)
return CameraConfig("render_camera", pose, 512, 512, 1, 0.01, 100)
[docs] def _load_agent(self, options: dict):
super()._load_agent(options, sapien.Pose(p=[-0.615, 0, 0]))
[docs] def _load_scene(self, options: dict):
self.table_scene = TableSceneBuilder(
self, robot_init_qpos_noise=self.robot_init_qpos_noise
)
self.table_scene.build()
self.cube = actors.build_cube(
self.scene,
half_size=self.cube_half_size,
color=[1, 0, 0, 1],
name="cube",
initial_pose=sapien.Pose(p=[0, 0, self.cube_half_size]),
)
self.goal_site = actors.build_sphere(
self.scene,
radius=self.goal_thresh,
color=[0, 1, 0, 1],
name="goal_site",
body_type="kinematic",
add_collision=False,
initial_pose=sapien.Pose(),
)
self._hidden_objects.append(self.goal_site)
[docs] def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
with torch.device(self.device):
b = len(env_idx)
self.table_scene.initialize(env_idx)
xyz = torch.zeros((b, 3))
xyz[:, :2] = (
torch.rand((b, 2)) * self.cube_spawn_half_size * 2
- self.cube_spawn_half_size
)
xyz[:, 0] += self.cube_spawn_center[0]
xyz[:, 1] += self.cube_spawn_center[1]
xyz[:, 2] = self.cube_half_size
qs = randomization.random_quaternions(b, lock_x=True, lock_y=True)
self.cube.set_pose(Pose.create_from_pq(xyz, qs))
goal_xyz = torch.zeros((b, 3))
goal_xyz[:, :2] = (
torch.rand((b, 2)) * self.cube_spawn_half_size * 2
- self.cube_spawn_half_size
)
goal_xyz[:, 0] += self.cube_spawn_center[0]
goal_xyz[:, 1] += self.cube_spawn_center[1]
goal_xyz[:, 2] = torch.rand((b)) * self.max_goal_height + xyz[:, 2]
self.goal_site.set_pose(Pose.create_from_pq(goal_xyz))
[docs] def evaluate(self):
is_obj_placed = (
torch.linalg.norm(self.goal_site.pose.p - self.cube.pose.p, axis=1)
<= self.goal_thresh
)
is_grasped = self.agent.is_grasping(self.cube)
is_robot_static = self.agent.is_static(0.2)
return {
"success": is_obj_placed & is_robot_static,
"is_obj_placed": is_obj_placed,
"is_robot_static": is_robot_static,
"is_grasped": is_grasped,
}
[docs] def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
tcp_to_obj_dist = torch.linalg.norm(
self.cube.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.cube.pose.p, axis=1
)
place_reward = 1 - torch.tanh(5 * obj_to_goal_dist)
reward += place_reward * is_grasped
qvel = self.agent.robot.get_qvel()
if self.robot_uids in ["panda", "widowxai"]:
qvel = qvel[..., :-2]
elif self.robot_uids == "so100":
qvel = qvel[..., :-1]
static_reward = 1 - torch.tanh(5 * torch.linalg.norm(qvel, axis=1))
reward += static_reward * info["is_obj_placed"]
reward[info["success"]] = 5
return reward
[docs] def compute_normalized_dense_reward(
self, obs: Any, action: torch.Tensor, info: dict
):
return self.compute_dense_reward(obs=obs, action=action, info=info) / 5
PickCubeEnv.__doc__ = PICK_CUBE_DOC_STRING.format(robot_id="Panda")
@register_env("PickCubeSO100-v1", max_episode_steps=50)
[docs]class PickCubeSO100Env(PickCubeEnv):
[docs] _sample_video_link = "https://github.com/haosulab/ManiSkill/raw/main/figures/environment_demos/PickCubeSO100-v1_rt.mp4"
def __init__(self, *args, **kwargs):
super().__init__(*args, robot_uids="so100", **kwargs)
PickCubeSO100Env.__doc__ = PICK_CUBE_DOC_STRING.format(robot_id="SO100")
@register_env("PickCubeWidowXAI-v1", max_episode_steps=50)
[docs]class PickCubeWidowXAIEnv(PickCubeEnv):
[docs] _sample_video_link = "https://github.com/haosulab/ManiSkill/raw/main/figures/environment_demos/PickCubeWidowXAI-v1_rt.mp4"
def __init__(self, *args, **kwargs):
super().__init__(*args, robot_uids="widowxai", **kwargs)
PickCubeWidowXAIEnv.__doc__ = PICK_CUBE_DOC_STRING.format(robot_id="WidowXAI")
