Skip to content

Commit 0d74b07

Browse files
peterdavidfaganpeterdavidfagan
committed
intermediary update
1 parent 1c32aa7 commit 0d74b07

File tree

2 files changed

+372
-0
lines changed

2 files changed

+372
-0
lines changed

‎mujoco_robot_environments/tasks/lasa_draw.py‎

Lines changed: 371 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,371 @@
1+
"""Mujoco environment for prototyping LASA dataset drawing motions."""
2+
from abc import abstractmethod
3+
from typing import Optional, Dict
4+
from copy import deepcopy
5+
from pathlib import Path
6+
import random
7+
8+
import mujoco
9+
from mujoco import viewer
10+
from mujoco import mj_name2id, mj_id2name
11+
import jax.numpy as jnp
12+
import numpy as np
13+
from scipy.spatial.transform import Rotation as R
14+
import dm_env
15+
from dm_control import composer, mjcf
16+
from dm_control.composer.variation import distributions
17+
from dm_control.composer.variation import rotations
18+
import hydra
19+
from hydra.utils import instantiate
20+
from hydra import compose, initialize
21+
from omegaconf import DictConfig
22+
23+
from mujoco_robot_environments.models.arenas import empty
24+
from mujoco_robot_environments.models.robot_arm import standard_compose
25+
from mujoco_robot_environments.environment.props import add_objects, Rectangle
26+
from mujoco_robot_environments.environment.cameras import add_camera
27+
from mujoco_robot_environments.environment.prop_initializer import PropPlacer
28+
from mujoco_robot_environments.models.robot_arm import RobotArm
29+
30+
31+
def generate_default_config():
32+
hydra.core.global_hydra.GlobalHydra.instance().clear()
33+
initialize(version_base=None, config_path="../config", job_name="rearrangement")
34+
return compose(
35+
config_name="rearrangement",
36+
overrides=[
37+
"arena/props=colour_splitter",
38+
"simulation_tuning_mode=False"
39+
]
40+
)
41+
42+
43+
class PushEnv(dm_env.Environment):
44+
"""MuJoCo powered robotics environment with dm_env interface."""
45+
46+
def __init__(
47+
self,
48+
viewer: Optional = None,
49+
cfg: DictConfig = generate_default_config(),
50+
):
51+
"""Initializes the simulation environment from config."""
52+
# ensure mjcf paths are relative to this file
53+
file_path = Path(__file__).parent.absolute()
54+
self._cfg = cfg
55+
56+
# check if viewer is requested in input args, otherwise use config
57+
if viewer is not None:
58+
self.has_viewer = viewer
59+
elif self._cfg.viewer is None:
60+
self.has_viewer = False
61+
print("Viewer not requested, running headless.")
62+
else:
63+
self.has_viewer = self._cfg.viewer
64+
65+
# create arena
66+
self._arena = empty.Arena()
67+
68+
# set general physics parameters
69+
self._arena.mjcf_model.option.timestep = cfg.physics_dt
70+
self._arena.mjcf_model.option.gravity = cfg.gravity
71+
self._arena.mjcf_model.size.nconmax = cfg.nconmax
72+
self._arena.mjcf_model.size.njmax = cfg.njmax
73+
self._arena.mjcf_model.visual.__getattr__("global").offheight = cfg.offheight
74+
self._arena.mjcf_model.visual.__getattr__("global").offwidth = cfg.offwidth
75+
self._arena.mjcf_model.visual.map.znear = cfg.znear
76+
77+
# add table for manipulation task
78+
table = Rectangle(
79+
name="table",
80+
x_len=1.0,
81+
y_len=1.0,
82+
z_len=0.2,
83+
rgba=(1.0, 1.0, 1.0, 1.0),
84+
margin=0.0,
85+
gap=0.0,
86+
mass=10,
87+
)
88+
89+
table_attach_site = self._arena.mjcf_model.worldbody.add(
90+
"site",
91+
name=f"table_center",
92+
pos=(0.4, 0.0, 0.2),
93+
)
94+
self._arena.attach(table, table_attach_site)
95+
96+
# add robot model with actuators and sensors
97+
self.arm = instantiate(cfg.robots.arm.arm)
98+
99+
# add cylinder to end effector for non-prehensile manipulation
100+
cylinder_geom = self.arm._attachment_site.parent.add(
101+
'geom',
102+
type='cylinder',
103+
size=[0.015, 0.05], # radius and half-height
104+
pos=[0.0, 0.0, 0.05],
105+
rgba=[0.02, 0.302, 0.4, 1.0] # red color
106+
)
107+
108+
robot_base_site = self._arena.mjcf_model.worldbody.add(
109+
"site",
110+
name="robot_base",
111+
pos=(0.0, 0.0, 0.4),
112+
)
113+
self._arena.attach(self.arm, robot_base_site)
114+
115+
116+
# if debugging the task environment add mocap for controller eef
117+
if cfg.simulation_tuning_mode:
118+
self.eef_target_mocap=self._arena.mjcf_model.worldbody.add(
119+
'body',
120+
name="eef_target_mocap",
121+
mocap="true",
122+
pos=[0.6, -0.25, 0.55],
123+
quat=R.from_euler('xyz', [180, 180, 0], degrees=True).as_quat(),
124+
)
125+
126+
self.eef_target_mocap.add(
127+
'geom',
128+
name='mocap_target_viz',
129+
type='box',
130+
size=[0.025, 0.025, 0.025],
131+
rgba=[1.0, 0.0, 0.0, 0.25],
132+
pos=[0.0, 0.0, 0.0],
133+
contype=0, # no collision with any object
134+
conaffinity=0 # no influence on collision detection
135+
)
136+
137+
# visualise plane for drawing
138+
self.draw_plane_center=self._arena.mjcf_model.worldbody.add(
139+
'body',
140+
name="draw_plane",
141+
mocap="false",
142+
pos=[0.4, 0.0, 0.55],
143+
quat=R.from_euler('xyz', [180, 180, 0], degrees=True).as_quat(),
144+
)
145+
146+
self.draw_plane_center.add(
147+
'geom',
148+
name='draw_plane',
149+
type='box',
150+
size=[0.3, 0.4, 0.001],
151+
rgba=[0.0, 0.0, 1.0, 0.05],
152+
pos=[0.0, 0.0, 0.0],
153+
contype=0, # no collision with any object
154+
conaffinity=0 # no influence on collision detection
155+
)
156+
157+
158+
# add cameras
159+
for camera in cfg.arena.cameras:
160+
add_camera(
161+
self._arena,
162+
name=camera.name,
163+
pos=camera.pos,
164+
quat=camera.quat,
165+
height=camera.height,
166+
width=camera.width,
167+
fovy=camera.fovy,
168+
)
169+
170+
# this environment uses this camera for observation specification
171+
if camera.name == "overhead_camera":
172+
self.overhead_camera_height = camera.height
173+
self.overhead_camera_width = camera.width
174+
175+
# compile environment
176+
self._physics = mjcf.Physics.from_mjcf_model(self._arena.mjcf_model)
177+
self.renderer = mujoco.Renderer(self._physics.model.ptr, height=self.overhead_camera_height, width=self.overhead_camera_width)
178+
self.seg_renderer = mujoco.Renderer(self._physics.model.ptr, height=self.overhead_camera_height, width=self.overhead_camera_width)
179+
self.seg_renderer.enable_segmentation_rendering()
180+
self.depth_renderer = mujoco.Renderer(self._physics.model.ptr, height=self.overhead_camera_height, width=self.overhead_camera_width)
181+
self.depth_renderer.enable_depth_rendering()
182+
self.passive_view = None
183+
184+
def close(self) -> None:
185+
if self.passive_view is not None:
186+
self.passive_view.close()
187+
188+
@property
189+
def model(self) -> mujoco.MjModel:
190+
return self.physics.model
191+
192+
@property
193+
def data(self) -> mujoco.MjData:
194+
return self.physics.data
195+
196+
def reset(self) -> dm_env.TimeStep:
197+
"""Resets the environment to an initial state and returns the first
198+
`TimeStep` of the new episode.
199+
"""
200+
# reset the lation instance
201+
self._physics.reset()
202+
203+
# reset arm to home position
204+
# Note: for other envs we may want random sampling of initial arm positions
205+
self.arm.set_joint_angles(self._physics, self.arm.named_configurations["home"])
206+
print(self.arm.named_configurations["home"])
207+
208+
# configure viewer
209+
if self.has_viewer:
210+
if self.passive_view is not None:
211+
self.passive_view.close()
212+
self.passive_view = viewer.launch_passive(self._physics.model.ptr, self._physics.data.ptr)
213+
214+
# create an instance of a robot interface for robot and controllers
215+
self._robot = RobotArm(
216+
arm=self.arm,
217+
physics=self._physics,
218+
passive_viewer=self.passive_view,
219+
)
220+
221+
# set the initial eef pose to home
222+
self.eef_home_pose = self._robot.eef_pose.copy()
223+
224+
return dm_env.TimeStep(
225+
step_type=dm_env.StepType.FIRST,
226+
reward=0.0,
227+
discount=0.0,
228+
observation=self._compute_observation(),
229+
)
230+
231+
def step(self, action_dict) -> dm_env.TimeStep:
232+
"""
233+
Updates the environment according to the action and returns a `TimeStep`.
234+
"""
235+
observation = self._compute_observation()
236+
237+
return dm_env.TimeStep(
238+
step_type=dm_env.StepType.MID,
239+
reward=0.0,
240+
discount=0.0,
241+
observation=observation,
242+
)
243+
244+
def observation_spec(self) -> dm_env.specs.Array:
245+
"""Returns the observation spec."""
246+
# get shape of overhead camera
247+
camera = self._arena.mjcf_model.find("camera", "overhead_camera/overhead_camera")
248+
camera_shape = self.overhead_camera_height, self.overhead_camera_width, 3
249+
return {
250+
"overhead_camera/depth": dm_env.specs.Array(shape=camera_shape[:-1], dtype=np.float32),
251+
"overhead_camera/rgb": dm_env.specs.Array(shape=camera_shape, dtype=np.float32),
252+
}
253+
254+
def action_spec(self) -> Dict[str, dm_env.specs.Array]:
255+
"""Returns the action spec."""
256+
return {
257+
"pose": dm_env.specs.Array(shape=(7,), dtype=np.float64), # [x, y, z, qx, qy, qz, qw]
258+
"pixel_coords": dm_env.specs.Array(shape=(2,), dtype=np.int64), # [u, v]
259+
"gripper_rot": dm_env.specs.Array(shape=(1,), dtype=np.float64),
260+
}
261+
262+
def _compute_observation(self) -> np.ndarray:
263+
"""Returns the observation."""
264+
# get overhead camera
265+
camera_id = mj_name2id(self._physics.model.ptr, mujoco.mjtObj.mjOBJ_CAMERA, "overhead_camera/overhead_camera")
266+
self.renderer.update_scene(self._physics.data.ptr, camera_id)
267+
self.depth_renderer.update_scene(self._physics.data.ptr, camera_id)
268+
269+
# get rgb data
270+
rgb = self.renderer.render()
271+
272+
# get depth data
273+
depth = self.depth_renderer.render()
274+
275+
# add to observation
276+
obs = {}
277+
obs["overhead_camera/rgb"] = rgb
278+
obs["overhead_camera/depth"] = depth
279+
280+
return obs
281+
282+
def interactive_tuning(self):
283+
"""
284+
Interactively control arm to tune simulation parameters.
285+
"""
286+
287+
# get difference between eef site and mocap body
288+
mocap_pos = self._physics.data.mocap_pos[0]
289+
mocap_quat = self._physics.data.mocap_quat[0]
290+
291+
# update control target
292+
self._robot.arm_controller.set_target(
293+
position=mocap_pos + [0.0, 0.0, 0.1],
294+
quat=mocap_quat,
295+
velocity=np.zeros(3),
296+
angular_velocity=np.zeros(3),
297+
)
298+
299+
control_command = self._robot.arm_controller.compute_control_output()
300+
301+
# step the simulation
302+
for _ in range(5):
303+
self._physics.set_control(control_command)
304+
self._physics.step()
305+
if self.passive_view is not None:
306+
self.passive_view.sync()
307+
308+
def test_board(self, target_position):
309+
"""
310+
Interactively control arm to tune simulation parameters.
311+
"""
312+
mocap_quat = self._physics.data.mocap_quat[0]
313+
314+
# update control target
315+
self._robot.arm_controller.set_target(
316+
position=target_position + [0.0, 0.0, 0.1],
317+
quat=mocap_quat,
318+
velocity=np.zeros(3),
319+
angular_velocity=np.zeros(3),
320+
)
321+
322+
control_command = self._robot.arm_controller.compute_control_output()
323+
324+
# step the simulation
325+
for _ in range(5):
326+
self._physics.set_control(control_command)
327+
self._physics.step()
328+
if self.passive_view is not None:
329+
self.passive_view.sync()
330+
331+
332+
if __name__=="__main__":
333+
# clear hydra global state to avoid conflicts with other hydra instances
334+
hydra.core.global_hydra.GlobalHydra.instance().clear()
335+
336+
# read hydra config
337+
initialize(version_base=None, config_path="../config", job_name="rearrangement")
338+
339+
# add task configs
340+
COLOR_SEPARATING_CONFIG = compose(
341+
config_name="rearrangement",
342+
overrides=[
343+
"arena/props=colour_splitter",
344+
"simulation_tuning_mode=True"
345+
]
346+
)
347+
348+
# instantiate color separation task
349+
env = PushEnv(viewer=True, cfg=COLOR_SEPARATING_CONFIG)
350+
positions = [
351+
np.array([0.6, 0.3, 0.55]),
352+
np.array([0.6, -0.3, 0.55]),
353+
np.array([0.3, -0.3, 0.55]),
354+
np.array([0.3, 0.3, 0.55]),
355+
np.array([0.6, -0.3, 0.55]),
356+
]
357+
358+
# interactive control of robot with mocap body
359+
_, _, _, obs = env.reset()
360+
for target in positions:
361+
while True:
362+
env.test_board(target)
363+
364+
# check if target is reached
365+
if env._robot.arm_controller.current_position_error() < 1e-3:
366+
break
367+
368+
# env.interactive_tuning()
369+
370+
env.close()
371+

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ tensorflow-cpu = {version="^2.14.0", markers = "sys_platform == 'linux'"}
6161
envlogger = {extras = ["tfds"], version = "^1.2", markers = "sys_platform == 'linux'"}
6262
rlds = {version="^0.1.7", markers = "sys_platform == 'linux'"}
6363
robot-descriptions = "^1.10.0"
64+
mink = "^0.0.5"
6465

6566
[tool.black]
6667
line-length = 120

0 commit comments

Comments
 (0)