import numpy as np
import sapien
import torch
from transforms3d.euler import euler2quat
from mani_skill.agents.robots.panda.panda_stick import PandaStick
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.registration import register_env
from mani_skill.utils.scene_builder.table.scene_builder import TableSceneBuilder
from mani_skill.utils.structs.actor import Actor
from mani_skill.utils.structs.pose import Pose
from mani_skill.utils.structs.types import SceneConfig, SimConfig
@register_env("TableTopFreeDraw-v1", max_episode_steps=1000)
[docs]class TableTopFreeDrawEnv(BaseEnv):
"""
**Task Description:**
Instantiates a table with a white canvas on it and a robot with a stick that draws red lines. This environment is primarily for a reference / for others to copy
to make their own drawing tasks.
**Randomizations:**
None
**Success Conditions:**
None
**Goal Specification:**
None
"""
[docs] _sample_video_link = "https://github.com/haosulab/ManiSkill/raw/main/figures/environment_demos/TableTopFreeDraw-v1_rt.mp4"
"""
The total "ink" available to use and draw with before you need to call env.reset. NOTE that on GPU simulation it is not recommended to have a very high value for this as it can slow down rendering
when too many objects are being rendered in many scenes.
"""
"""thickness of the paint drawn on to the canvas"""
[docs] CANVAS_THICKNESS = 0.02
"""How thick the canvas on the table is"""
"""The brushes radius"""
[docs] BRUSH_COLORS = [[0.8, 0.2, 0.2, 1]]
"""The colors of the brushes. If there is more than one color, each parallel environment will have a randomly sampled color."""
[docs] SUPPORTED_REWARD_MODES = ["none"]
[docs] SUPPORTED_ROBOTS: ["panda_stick"]
def __init__(self, *args, robot_uids="panda_stick", **kwargs):
super().__init__(*args, robot_uids=robot_uids, **kwargs)
@property
[docs] def _default_sim_config(self):
# we set contact_offset to a small value as we are not expecting to make any contacts really apart from the brush hitting the canvas too hard.
# We set solver iterations very low as this environment is not doing a ton of manipulation (the brush is attached to the robot after all)
return SimConfig(
sim_freq=100,
control_freq=20,
scene_config=SceneConfig(
contact_offset=0.01,
solver_position_iterations=4,
solver_velocity_iterations=0,
),
)
@property
[docs] def _default_sensor_configs(self):
pose = sapien_utils.look_at(eye=[0.3, 0, 0.8], target=[0, 0, 0.1])
return [
CameraConfig(
"base_camera",
pose=pose,
width=320,
height=240,
fov=1.2,
near=0.01,
far=100,
)
]
@property
[docs] def _default_human_render_camera_configs(self):
pose = sapien_utils.look_at(eye=[0.3, 0, 0.8], target=[0, 0, 0.1])
return CameraConfig(
"render_camera",
pose=pose,
width=1280,
height=960,
fov=1.2,
near=0.01,
far=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=0)
self.table_scene.build()
# build a white canvas on the table
self.canvas = self.scene.create_actor_builder()
self.canvas.add_box_visual(
half_size=[0.4, 0.6, self.CANVAS_THICKNESS / 2],
material=sapien.render.RenderMaterial(base_color=[1, 1, 1, 1]),
)
self.canvas.add_box_collision(
half_size=[0.4, 0.6, self.CANVAS_THICKNESS / 2],
)
self.canvas.initial_pose = sapien.Pose(p=[-0.1, 0, self.CANVAS_THICKNESS / 2])
self.canvas = self.canvas.build_static(name="canvas")
self.dots = []
color_choices = torch.randint(0, len(self.BRUSH_COLORS), (self.num_envs,))
for i in range(self.MAX_DOTS):
actors = []
if len(self.BRUSH_COLORS) > 1:
for env_idx in range(self.num_envs):
builder = self.scene.create_actor_builder()
builder.add_cylinder_visual(
radius=self.BRUSH_RADIUS,
half_length=self.DOT_THICKNESS / 2,
material=sapien.render.RenderMaterial(
base_color=self.BRUSH_COLORS[color_choices[env_idx]]
),
)
builder.set_scene_idxs([env_idx])
builder.initial_pose = sapien.Pose(p=[0, 0, 0])
actor = builder.build_kinematic(name=f"dot_{i}_{env_idx}")
actors.append(actor)
self.dots.append(Actor.merge(actors))
else:
builder = self.scene.create_actor_builder()
builder.add_cylinder_visual(
radius=self.BRUSH_RADIUS,
half_length=self.DOT_THICKNESS / 2,
material=sapien.render.RenderMaterial(
base_color=self.BRUSH_COLORS[0]
),
)
builder.initial_pose = sapien.Pose(p=[0, 0, 0])
actor = builder.build_kinematic(name=f"dot_{i}")
self.dots.append(actor)
[docs] def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
# NOTE (stao): for simplicity this task cannot handle partial resets
self.draw_step = 0
with torch.device(self.device):
self.table_scene.initialize(env_idx)
for dot in self.dots:
# initially spawn dots in the table so they aren't seen
dot.set_pose(
sapien.Pose(
p=[0, 0, -self.DOT_THICKNESS], q=euler2quat(0, np.pi / 2, 0)
)
)
[docs] def _after_control_step(self):
if self.gpu_sim_enabled:
self.scene._gpu_fetch_all()
# This is the actual, GPU parallelized, drawing code.
# This is not real drawing but seeks to mimic drawing by placing dots on the canvas whenever the robot is close enough to the canvas surface
# We do not actually check if the robot contacts the table (although that is possible) and instead use a fast method to check.
# We add a 0.005 meter of leeway to make it easier for the robot to get close to the canvas and start drawing instead of having to be super close to the table.
robot_touching_table = (
self.agent.tcp.pose.p[:, 2]
< self.CANVAS_THICKNESS + self.DOT_THICKNESS + 0.005
)
robot_brush_pos = torch.zeros((self.num_envs, 3), device=self.device)
robot_brush_pos[:, 2] = -self.DOT_THICKNESS
robot_brush_pos[robot_touching_table, :2] = self.agent.tcp.pose.p[
robot_touching_table, :2
]
robot_brush_pos[robot_touching_table, 2] = (
self.DOT_THICKNESS / 2 + self.CANVAS_THICKNESS
)
# move the next unused dot to the robot's brush position. All unused dots are initialized inside the table so they aren't visible
self.dots[self.draw_step].set_pose(
Pose.create_from_pq(robot_brush_pos, euler2quat(0, np.pi / 2, 0))
)
self.draw_step += 1
# on GPU sim we have to call _gpu_apply_all() to apply the changes we make to object poses.
if self.gpu_sim_enabled:
self.scene._gpu_apply_all()
[docs] def evaluate(self):
return {}
