from typing import Union
import numpy as np
import sapien
import torch
from mani_skill.agents.robots import Fetch, Panda
from mani_skill.envs.sapien_env import BaseEnv
from mani_skill.envs.utils import randomization
from mani_skill.sensors.camera import CameraConfig
from mani_skill.utils import common, sapien_utils
from mani_skill.utils.building import actors
from mani_skill.utils.geometry.rotation_conversions import (
euler_angles_to_matrix,
matrix_to_quaternion,
)
from mani_skill.utils.logging_utils import logger
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
@register_env("StackPyramid-v1", max_episode_steps=250)
[docs]class StackPyramidEnv(BaseEnv):
"""
**Task Description:**
- The goal is to pick up a red cube, place it next to the green cube, and stack the blue cube on top of the red and green cube without it falling off.
**Randomizations:**
- all cubes have their z-axis rotation randomized
- all cubes have their xy positions on top of the table scene randomized. The positions are sampled such that the cubes do not collide with each other
**Success Conditions:**
- the blue cube is static
- the blue cube is on top of both the red and green cube (to within half of the cube size)
- none of the red, green, blue cubes are grasped by the robot (robot must let go of the cubes)
_sample_video_link = "https://github.com/haosulab/ManiSkill/raw/main/figures/environment_demos/StackPyramid-v1_rt.mp4"
"""
[docs] SUPPORTED_ROBOTS = ["panda_wristcam", "panda", "fetch"]
[docs] SUPPORTED_REWARD_MODES = ["none", "sparse"]
[docs] agent: Union[Panda, Fetch]
def __init__(
self, *args, robot_uids="panda_wristcam", robot_init_qpos_noise=0.02, **kwargs
):
[docs] self.robot_init_qpos_noise = robot_init_qpos_noise
super().__init__(*args, robot_uids=robot_uids, **kwargs)
@property
[docs] def _default_sensor_configs(self):
pose = sapien_utils.look_at(eye=[0.3, 0, 0.4], target=[-0.05, 0, 0.1])
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([0.6, 0.7, 0.6], [0.0, 0.0, 0.35])
return CameraConfig("render_camera", pose, 512, 512, 1, 0.01, 100)
[docs] def _load_scene(self, options: dict):
self.cube_half_size = common.to_tensor([0.02] * 3)
self.table_scene = TableSceneBuilder(
env=self, robot_init_qpos_noise=self.robot_init_qpos_noise
)
self.table_scene.build()
self.cubeA = actors.build_cube(
self.scene,
half_size=0.02,
color=[1, 0, 0, 1],
name="cubeA",
initial_pose=sapien.Pose(p=[0, 0, 0.2]),
)
self.cubeB = actors.build_cube(
self.scene,
half_size=0.02,
color=[0, 1, 0, 1],
name="cubeB",
initial_pose=sapien.Pose(p=[1, 0, 0.2]),
)
self.cubeC = actors.build_cube(
self.scene,
half_size=0.02,
color=[0, 0, 1, 1],
name="cubeC",
initial_pose=sapien.Pose(p=[-1, 0, 0.2]),
)
[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), device=self.device)
xyz[:, 2] = 0.02
xy = xyz[:, :2]
region = [[-0.1, -0.2], [0.1, 0.2]]
sampler = randomization.UniformPlacementSampler(
bounds=region, batch_size=b, device=self.device
)
radius = torch.linalg.norm(torch.tensor([0.02, 0.02]))
cubeA_xy = xy + sampler.sample(radius, 100)
cubeB_xy = xy + sampler.sample(radius, 100, verbose=False)
cubeC_xy = xy + sampler.sample(radius, 100, verbose=False)
# Cube A
xyz[:, :2] = cubeA_xy
qs = randomization.random_quaternions(
b,
lock_x=True,
lock_y=True,
lock_z=False,
)
self.cubeA.set_pose(Pose.create_from_pq(p=xyz.clone(), q=qs))
# Cube B
xyz[:, :2] = cubeB_xy
qs = randomization.random_quaternions(
b,
lock_x=True,
lock_y=True,
lock_z=False,
)
self.cubeB.set_pose(Pose.create_from_pq(p=xyz.clone(), q=qs))
# Cube C
xyz[:, :2] = cubeC_xy
qs = randomization.random_quaternions(
b,
lock_x=True,
lock_y=True,
lock_z=False,
)
self.cubeC.set_pose(Pose.create_from_pq(p=xyz, q=qs))
[docs] def evaluate(self):
pos_A = self.cubeA.pose.p
pos_B = self.cubeB.pose.p
pos_C = self.cubeC.pose.p
offset_AB = pos_A - pos_B
offset_BC = pos_B - pos_C
offset_AC = pos_A - pos_C
def evaluate_cube_distance(offset, cube_a, cube_b, top_or_next):
xy_flag = (
torch.linalg.norm(offset[..., :2], axis=1)
<= torch.linalg.norm(2 * self.cube_half_size[:2]) + 0.005
)
z_flag = torch.abs(offset[..., 2]) > 0.02
if top_or_next == "top":
is_cubeA_on_cubeB = torch.logical_and(xy_flag, z_flag)
elif top_or_next == "next_to":
is_cubeA_on_cubeB = xy_flag
else:
return NotImplementedError(
f"Expect top_or_next to be either 'top' or 'next_to', got {top_or_next}"
)
is_cubeA_static = cube_a.is_static(lin_thresh=1e-2, ang_thresh=0.5)
is_cubeA_grasped = self.agent.is_grasping(cube_a)
success = is_cubeA_on_cubeB & is_cubeA_static & (~is_cubeA_grasped)
return success.bool()
success_A_B = evaluate_cube_distance(
offset_AB, self.cubeA, self.cubeB, "next_to"
)
success_C_B = evaluate_cube_distance(offset_BC, self.cubeC, self.cubeB, "top")
success_C_A = evaluate_cube_distance(offset_AC, self.cubeC, self.cubeA, "top")
success = torch.logical_and(
success_A_B, torch.logical_and(success_C_B, success_C_A)
)
return {
"success": success,
}
