def test_multiple_sims():
# Ensure that creating new simulators still produces good renderings.
xml = """
<mujoco>
<asset>
<texture name="t1" width="32" height="32" type="2d" builtin="flat" />
<material name="m1" texture="t1" />
</asset>
<worldbody>
<light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
<camera name="topcam" pos="0 0 2.5" zaxis="0 0 1" />
<geom name="g1" pos="0 0 0" type="box" size="1 1 0.1" rgba="1 1 1 1" material="m1" />
</worldbody>
</mujoco>
"""
model = load_model_from_xml(xml)
random_state = np.random.RandomState(0)
for i in range(3):
sim = MjSim(model)
sim.forward()
modder = TextureModder(sim, random_state=random_state)
for j in range(2):
modder.rand_checker('g1')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_multiple_sims.loop%d_%d.png' % (i, j))
def test():
from mujoco_py import load_model_from_xml, MjSim, MjViewer
from mujoco_py.modder import TextureModder
l, w, h, t, left, m = sample_cabinet()
cab = build_cabinet(l, w, h, t, left)
filename = 'test.urdf'
with open(filename, "w") as text_file:
text_file.write(cab.xml)
# print(cab.xml)
model = load_model_from_xml(cab.xml)
sim = MjSim(model)
viewer = MjViewer(sim)
modder = TextureModder(sim)
for name in sim.model.geom_names:
modder.rand_all(name)
t = 0
if cab.geom[3] == 1:
sim.data.ctrl[0] = -0.2
else:
sim.data.ctrl[0] = 0.2
while t < 1000:
# for name in sim.model.geom_names:
# modder.rand_all(name)
sim.step()
viewer.render()
t += 1
def test_resetting():
# Ensure that resetting environment and creating new simulators
# still produces good renderings.
xml = """
<mujoco>
<asset>
<texture name="t1" width="32" height="32" type="2d" builtin="flat" />
<material name="m1" texture="t1" />
</asset>
<worldbody>
<light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
<camera name="topcam" pos="0 0 2.5" zaxis="0 0 1" />
<geom name="g1" pos="0 0 0" type="{geom_type}" size="1 1 0.1" rgba="1 1 1 1" material="m1" />
</worldbody>
</mujoco>
"""
def get_sim(seed):
geom_type = ["box", "sphere"][seed % 2]
model = load_model_from_xml(xml.format(geom_type=geom_type))
return MjSim(model)
random_state = np.random.RandomState(0)
for i in range(3):
sim = get_sim(i - 1)
sim.forward()
modder = TextureModder(sim, random_state=random_state)
for j in range(2):
modder.rand_checker('g1')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_resetting.loop%d_%d.png' % (i, j))
def test_multiple_sims():
# Ensure that creating new simulators still produces good renderings.
xml = """
<mujoco>
<asset>
<texture name="t1" width="32" height="32" type="2d" builtin="flat" />
<material name="m1" texture="t1" />
</asset>
<worldbody>
<light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
<camera name="topcam" pos="0 0 2.5" zaxis="0 0 1" />
<geom name="g1" pos="0 0 0" type="box" size="1 1 0.1" rgba="1 1 1 1" material="m1" />
</worldbody>
</mujoco>
"""
model = load_model_from_xml(xml)
random_state = np.random.RandomState(0)
for i in range(3):
sim = MjSim(model)
sim.forward()
modder = TextureModder(sim, random_state=random_state)
for j in range(2):
modder.rand_checker('g1')
compare_imgs(
sim.render(201, 205, camera_name="topcam"),
'test_multiple_sims.loop%d_%d.png' % (i, j))
def test_resetting():
# Ensure that resetting environment and creating new simulators
# still produces good renderings.
xml = """
<mujoco>
<asset>
<texture name="t1" width="32" height="32" type="2d" builtin="flat" />
<material name="m1" texture="t1" />
</asset>
<worldbody>
<light diffuse=".5 .5 .5" pos="0 0 5" dir="0 0 -1" />
<camera name="topcam" pos="0 0 2.5" zaxis="0 0 1" />
<geom name="g1" pos="0 0 0" type="{geom_type}" size="1 1 0.1" rgba="1 1 1 1" material="m1" />
</worldbody>
</mujoco>
"""
def get_sim(seed):
geom_type = ["box", "sphere"][seed % 2]
model = load_model_from_xml(xml.format(geom_type=geom_type))
return MjSim(model)
random_state = np.random.RandomState(0)
for i in range(3):
sim = get_sim(i - 1)
sim.forward()
modder = TextureModder(sim, random_state=random_state)
for j in range(2):
modder.rand_checker('g1')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_resetting.loop%d_%d.png' % (i, j))
def test(k=0):
from mujoco_py import load_model_from_xml, MjSim, MjViewer
from mujoco_py.modder import TextureModder
l, w, h, t, left, m = sample_cabinet2()
cab = build_cabinet2(l, w, h, t, left)
# print(cab.xml)
model = load_model_from_xml(cab.xml)
sim = MjSim(model)
viewer = MjViewer(sim)
modder = TextureModder(sim)
for name in sim.model.geom_names:
modder.rand_all(name)
set_two_door_control(sim)
q1s = []
q2s = []
t = 0
# mode 0 indicates starting lc
while t < 4000:
# for name in sim.model.geom_names:
# modder.rand_all(name)
viewer.render()
if t % 250 == 0:
q1 = sim.data.qpos[0]
q2 = sim.data.qpos[1]
print(sim.data.qpos)
q1s.append(q1)
q2s.append(q2)
sim.step()
t += 1
# print(q1s)
np.save('devdata/q1_' + str(k).zfill(3), q1s)
np.save('devdata/q2_' + str(k).zfill(3), q2s)
def _env_setup(self, initial_qpos):
self.id2obj = [self._geom2objid(i) for i in range(self.sim.model.ngeom)]
for name, value in initial_qpos.items():
self.sim.data.set_joint_qpos(name, value)
utils.reset_mocap_welds(self.sim)
self.sim.forward()
# Move end effector into position.
gripper_target = np.array([-0.498, 0.005, -0.431 + self.gripper_extra_height]) + self.sim.data.get_site_xpos('robot0:grip')
gripper_rotation = np.array([1., 0., 1., 0.])
self.sim.data.set_mocap_pos('robot0:mocap', gripper_target)
self.sim.data.set_mocap_quat('robot0:mocap', gripper_rotation)
for _ in range(10):
self.sim.step()
# Extract information for sampling goals.
if not hasattr(self, 'initial_gripper_xpos'):
self.initial_gripper_xpos = self.sim.data.get_site_xpos('robot0:grip').copy()
self.height_offset = self.sim.data.get_site_xpos('object0')[2]
# Change the colors to match the success conditions
self.color_modder = TextureModder(self.sim)
for i in range(self.sim.model.ngeom):
obj_id = self.id2obj[i]
if obj_id != None:
name = self.sim.model.geom_id2name(i)
if 'table' in name:
color = np.asarray(COLORS_RGB[self.obj_colors[obj_id]])
self.color_modder.set_rgb(name, color * 2)
else:
self.color_modder.set_rgb(name,
COLORS_RGB[self.obj_colors[obj_id]])
def __init__(self,
model_path,
dataset_name,
use_procedural=False,
cam_pos_file=None,
cam_norm_pos_file=None):
super().__init__(dataset_name,
cam_pos_file=cam_pos_file,
cam_norm_pos_file=cam_norm_pos_file)
self.model = mujoco_py.load_model_from_path(model_path)
self.sim = mujoco_py.MjSim(self.model)
self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, None)
self.cam_modder = CameraModder(self.sim)
self.tex_modder = TextureModder(self.sim)
self.use_procedural = use_procedural
def _re_init(self, xml):
# TODO: Now, likely needs rank
randomized_path = os.path.join(self.xml_dir, "randomizedgen3.xml")
with open(randomized_path, 'wb') as fp:
fp.write(xml.encode())
fp.flush()
try:
self.model = mujoco_py.load_model_from_path(randomized_path)
except:
print("Unable to load the xml file")
self.sim = mujoco_py.MjSim(self.model, nsubsteps=self.n_substeps)
self.modder = TextureModder(self.sim)
self.mat_modder = MaterialModder(self.sim)
self.light_modder = LightModder(self.sim)
self.camera_modder = CameraModder(self.sim)
self.metadata = {
'render.modes': ['human', 'rgb_array'],
'video.frames_per_second': int(np.round(1.0 / self.dt))
}
self._env_setup(initial_qpos=self.initial_qpos)
self.initial_state = copy.deepcopy(self.sim.get_state())
if self.viewer:
self.viewer.update_sim(self.sim)
def test():
import cv2
from mujoco_py import load_model_from_xml, MjSim, MjViewer
from mujoco_py.modder import TextureModder
for i in range(100):
l, w, h, t, left, m = sample_refrigerator(False)
fridge = build_refrigerator(l,
w,
h,
t,
left,
set_pose=(3.0, 0.0, -1.0),
set_rot=(0, 0, 0, 1))
model = load_model_from_xml(fridge.xml)
sim = MjSim(model)
viewer = MjViewer(sim)
modder = TextureModder(sim)
set_two_door_control(sim, 'refrigerator')
t = 0
while t < 2000:
sim.step()
viewer.render()
t += 1
def __init__(self, model_path, initial_qpos, n_actions, n_substeps):
if model_path.startswith('/'):
fullpath = model_path
else:
fullpath = os.path.join(os.path.dirname(__file__), 'assets',
model_path)
if not os.path.exists(fullpath):
raise IOError('File {} does not exist'.format(fullpath))
model = mujoco_py.load_model_from_path(fullpath)
self.sim = mujoco_py.MjSim(model, nsubsteps=n_substeps)
#self.viewer = None # comment when using "human"
self.viewer = mujoco_py.MjViewer(
self.sim) #comment when using "rgb_array"
self._viewers = {}
self.modder = TextureModder(self.sim)
self.visual_randomize = True
self.mat_modder = MaterialModder(self.sim)
self.light_modder = LightModder(self.sim)
self.metadata = {
'render.modes': ['human', 'rgb_array'],
'video.frames_per_second': int(np.round(1.0 / self.dt))
}
self.visual_data_recording = True
self._index = 0
self._label_matrix = []
self.seed()
self._env_setup(initial_qpos=initial_qpos)
self.initial_state = copy.deepcopy(self.sim.get_state())
self.goal = self._sample_goal()
obs = self._get_obs()
self.action_space = spaces.Box(-1.,
1.,
shape=(n_actions, ),
dtype='float32')
self.observation_space = spaces.Dict(
dict(
desired_goal=spaces.Box(-np.inf,
np.inf,
shape=obs['achieved_goal'].shape,
dtype='float32'),
achieved_goal=spaces.Box(-np.inf,
np.inf,
shape=obs['achieved_goal'].shape,
dtype='float32'),
observation=spaces.Box(-np.inf,
np.inf,
shape=obs['observation'].shape,
dtype='float32'),
))
if self.viewer:
self._viewer_setup()
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# Get start state of params to slightly jitter later
self.start_geo_size = self.model.geom_size.copy()
self.start_geom_quat = self.model.geom_quat.copy()
self.start_body_pos = self.model.body_pos.copy()
self.start_body_quat = self.model.body_quat.copy()
self.start_matid = self.model.geom_matid.copy()
self.floor_offset = self.model.body_pos[self.model.body_name2id(
'floor')]
self.rock_mod_cache = None
self.tex_modder = TextureModder(self.sim)
self.cam_modder = CameraModder(self.sim)
self.light_modder = LightModder(self.sim)
# Set these both externally to visualize
self.visualize = False
self.viewer = None
def test():
# import transforms3d as tf3d
from mujoco_py import load_model_from_xml, MjSim, MjViewer
from mujoco_py.modder import TextureModder
for i in range(200):
l, w, h, t, left, m = sample_drawers()
drawer = build_drawer(l, w, h, t, left)
model = load_model_from_xml(drawer.xml)
sim = MjSim(model)
viewer = MjViewer(sim)
modder = TextureModder(sim)
for name in sim.model.geom_names:
modder.rand_all(name)
t = 0
sim.data.ctrl[0] = 0.05
while t < 1000:
sim.step()
viewer.render()
t += 1
def test():
from mujoco_py import load_model_from_xml, MjSim, MjViewer
from mujoco_py.modder import TextureModder
l, w, h, t, left, m = sample_toaster()
cab = build_toaster(l, w, h, t, left)
# print(cab.xml)
model = load_model_from_xml(cab.xml)
sim = MjSim(model)
viewer = MjViewer(sim)
modder = TextureModder(sim)
for name in sim.model.geom_names:
modder.rand_all(name)
t = 0
sim.data.ctrl[0] = 0.2
while t < 1000:
# for name in sim.model.geom_names:
# modder.rand_all(name)
sim.step()
viewer.render()
t += 1
def __init__(self,
env_file,
random_color=None,
random_num=5,
debug=False,
random_seed=0):
self.env_file = env_file
self.random_color = random_color
self.random_num = random_num
self.debug = debug
self.model = load_model_from_path(env_file)
self.sim = MjSim(self.model, nsubsteps=1)
self.sim.model.vis.map.znear = 0.02
self.sim.model.vis.map.zfar = 50.0
self.cube_size = self.sim.model.geom_size[
self.model._geom_name2id['cube_0']]
self.cuboid_size = self.sim.model.geom_size[
self.model._geom_name2id['cuboid_0']]
self.cube_num = len(
[i for i in self.sim.model.geom_names if 'cube_' in i])
self.cuboid_num = len(
[i for i in self.sim.model.geom_names if 'cuboid_' in i])
if self.debug:
print('total cube num:', self.cube_num)
print('total cuboid num:', self.cuboid_num)
self.max_num_per_type = max(self.cube_num, self.cuboid_num)
self.center_bounds = [
-0.25, 0.25
] # [0, self.cuboid_size[0] * self.max_num_per_type]
self.pos_candidates = np.arange(
self.center_bounds[0], self.center_bounds[1] + self.cube_size[0],
self.cube_size[0])
self.modder = TextureModder(self.sim, random_state=random_seed)
self.cur_id = {'cube': 0, 'cuboid': 0}
self.viewer = None
self.active_blocks = []
if random_color:
self.reset_viewer()
def __init__(self, model_path, initial_qpos, n_actions, n_substeps):
if model_path.startswith('/'):
fullpath = model_path
else:
fullpath = os.path.join(os.path.dirname(__file__), 'assets',
model_path)
if not os.path.exists(fullpath):
raise IOError('File {} does not exist'.format(fullpath))
model = mujoco_py.load_model_from_path(fullpath)
self.sim = mujoco_py.MjSim(model, nsubsteps=n_substeps)
self.modder = TextureModder(self.sim)
self.cam_modder = CameraModder(self.sim)
self.light_modder = LightModder(self.sim)
self.viewer = None
self.metadata = {
'render.modes': ['human', 'rgb_array'],
'video.frames_per_second': int(np.round(1.0 / self.dt))
}
self.seed()
self._env_setup(initial_qpos=initial_qpos)
self.initial_state = copy.deepcopy(self.sim.get_state())
self.goal = self._sample_goal()
obs = self._get_obs()
self.action_space = spaces.Box(-1.,
1.,
shape=(n_actions, ),
dtype='float32')
self.observation_space = spaces.Dict(
dict(
desired_goal=spaces.Box(-np.inf,
np.inf,
shape=obs['achieved_goal'].shape,
dtype='float32'),
achieved_goal=spaces.Box(-np.inf,
np.inf,
shape=obs['achieved_goal'].shape,
dtype='float32'),
observation=spaces.Box(-np.inf,
np.inf,
shape=obs['observation'].shape,
dtype='float32'),
))
self.episodeAcs = []
self.episodeObs = []
self.episodeInfo = []
def visualize_prediction(scene_filename, gt_axes, pred_axes):
tree = ET.parse(scene_filename)
root = tree.getroot()
# Load scene in mujoco
model = load_model_from_path(scene_filename)
sim = MjSim(model)
modder = TextureModder(sim)
# viewer=MjViewer(sim)
# Plot GT axis and predicted axis
# Take snapshot
print("Should have saved a snapshot")
def __init__(self, env_params_dict, reset_state=None):
hand_low = (-0.2, 0.4, 0.0)
hand_high = (0.2, 0.8, 0.05)
obj_low = (-0.3, 0.4, 0.1)
obj_high = (0.3, 0.8, 0.3)
dirname = '/'.join(os.path.abspath(__file__).split('/')[:-1])
params_dict = copy.deepcopy(env_params_dict)
_hp = self._default_hparams()
for name, value in params_dict.items():
print('setting param {} to value {}'.format(name, value))
_hp.set_hparam(name, value)
if _hp.textured:
filename = os.path.join(
dirname, "assets/sawyer_xyz/sawyer_multiobject_textured.xml")
else:
filename = os.path.join(
dirname, "assets/sawyer_xyz/sawyer_multiobject.xml")
BaseMujocoEnv.__init__(self, filename, _hp)
SawyerXYZEnv.__init__(self,
frame_skip=5,
action_scale=1. / 10,
hand_low=hand_low,
hand_high=hand_high,
model_name=filename)
if _hp.randomize_object_colors:
from mujoco_py.modder import TextureModder
self.texture_modder = TextureModder(self.sim)
goal_low = self.hand_low
goal_high = self.hand_high
self._adim = 4
self._hp = _hp
self.liftThresh = 0.04
self.max_path_length = 100
self.hand_init_pos = np.array((0, 0.6, 0.0))
def _re_init(self, xml):
# TODO: Now, likely needs rank
randomized_path = os.path.join(self.xml_dir, "randomizedgen3.xml")
with open(randomized_path, 'wb') as fp:
fp.write(xml.encode())
fp.flush()
try:
self.model = mujoco_py.load_model_from_path(randomized_path)
except:
print("Unable to load the xml file")
self.sim = mujoco_py.MjSim(self.model, nsubsteps=self.n_substeps)
self.modder = TextureModder(self.sim)
self.visual_randomize = True
self.metadata = {
'render.modes': ['human', 'rgb_array'],
'video.frames_per_second': int(np.round(1.0 / self.dt))
}
self._env_setup(initial_qpos=self.initial_qpos)
self.initial_state = copy.deepcopy(self.sim.get_state())
# self.initial_state = copy.deepcopy(self.sim.get_state())
# self.data = self.sim.data
# self.init_qpos = self.data.qpos.ravel().copy()
# self.init_qvel = self.data.qvel.ravel().copy()
# # observation = self.reset()
# #observation, _reward, done, _info = self.step(np.zeros(4))
# #assert not done
if self.viewer:
self.viewer.update_sim(self.sim)
#!/usr/bin/env python3
"""
Displays robot fetch at a disco party.
"""
from mujoco_py import load_model_from_path, MjSim, MjViewer
from mujoco_py.modder import TextureModder
import os
model = load_model_from_path("../xmls/fetch/main.xml")
pool = MjSimPool(model)
modder = TextureModder(sim)
pool = MjSimPool([MjSim(model) for _ in range(20)])
for i, sim in enumerate(pool.sims):
sim.data.qpos[:] = 0.0
sim.data.qvel[:] = 0.0
sim.data.ctrl[:] = i
# Advance all 20 simulations 100 times.
for _ in range(100):
pool.step()
import ipdb; ipdb.set_trace()
for i, sim in enumerate(pool.sims):
print("%d-th sim qpos=%s" % (i, str(sim.data.qpos)))
while True:
for name in sim.model.geom_names:
modder.rand_all(name)
class ArenaModder(BaseModder):
"""
Object to handle randomization of all relevant properties of Mujoco sim
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# Get start state of params to slightly jitter later
self.start_geo_size = self.model.geom_size.copy()
self.start_geom_quat = self.model.geom_quat.copy()
self.start_body_pos = self.model.body_pos.copy()
self.start_body_quat = self.model.body_quat.copy()
self.start_matid = self.model.geom_matid.copy()
self.floor_offset = self.model.body_pos[self.model.body_name2id(
'floor')]
self.rock_mod_cache = None
self.tex_modder = TextureModder(self.sim)
self.cam_modder = CameraModder(self.sim)
self.light_modder = LightModder(self.sim)
# Set these both externally to visualize
self.visualize = False
self.viewer = None
def whiten_materials(self):
self.tex_modder.whiten_materials()
def randomize(self):
self.mod_textures()
self.mod_lights()
self.mod_camera()
self.mod_walls()
self.mod_extras()
self.mod_rocks()
def mod_textures(self):
"""Randomize all the textures in the scene, including the skybox"""
for name in self.sim.model.geom_names:
if name != "billboard" and "light_disc" not in name:
self.tex_modder.rand_all(name)
##texture_sizes = [32, 64, 128, 256, 512, 1024]
##gid = self.model.geom_name2id(name)
##mid = self.model.geom_matid[gid]
##tid = self.model.mat_texid[mid]
##self.model.tex_width[tid] = sample_from_list(texture_sizes)
##self.model.tex_height[tid] = 6 * self.model.tex_width[tid]
self.tex_modder.rand_all('skybox')
def mod_lights(self):
"""Randomize pos, direction, and lights"""
# light stuff
LIGHT_RX = Range(LEFTX, RIGHTX)
LIGHT_RY = Range(BINY, DIGY)
LIGHT_RZ = Range(AFZ, AFZ + ZHIGH)
LIGHT_RANGE3D = Range3D(LIGHT_RX, LIGHT_RY, LIGHT_RZ)
LIGHT_UNIF = Range3D(Range(0, 1), Range(0, 1), Range(0, 1))
for i, name in enumerate(self.model.light_names):
lid = self.model.light_name2id(name)
# random sample 80% of any given light being on
self.light_modder.set_active(name, sample([0, 1]) < 0.8)
#self.light_modder.set_active(name, 0)
dir_xyz = sample_light_dir()
self.light_modder.set_pos(name, sample_xyz(LIGHT_RANGE3D))
self.light_modder.set_dir(name, dir_xyz)
self.light_modder.set_specular(name, sample_xyz(LIGHT_UNIF))
#self.light_modder.set_diffuse(name, sample_xyz(LIGHT_UNIF))
#self.light_modder.set_ambient(name, sample_xyz(LIGHT_UNIF))
#self.model.light_directional[lid] = sample([0,1]) < 0.01
def mod_camera(self):
"""Randomize pos, direction, and fov of camera"""
# Params
XOFF = 1.0
CAM_RX = Range(ACX - XOFF, ACX + XOFF) # center of arena +/- 0.5
CAM_RY = Range(BINY + 0.2, SZ_ENDY)
CAM_RZ = Range(AFZ + ZLOW, AFZ + ZHIGH)
CAM_RANGE3D = Range3D(CAM_RX, CAM_RY, CAM_RZ)
CAM_RYAW = Range(-95, -85)
CAM_RPITCH = Range(65, 90)
CAM_RROLL = Range(85, 95) # this might actually be pitch?
CAM_ANGLE3 = Range3D(CAM_RYAW, CAM_RPITCH, CAM_RROLL)
# "The horizontal field of view is computed automatically given the
# window size and the vertical field of view." - Mujoco
# This range was calculated using: themetalmuncher.github.io/fov-calc/
# ZED has 110° hfov --> 78° vfov, Logitech C920 has 78° hfov ---> 49° vfov
# These were rounded all the way down to 40° and up to 80°. It starts
# to look a bit bad in the upper range, but I think it will help
# generalization.
CAM_RFOVY = Range(40, 80)
# Actual mods
self.cam_modder.set_pos('camera1', sample_xyz(CAM_RANGE3D))
self.cam_modder.set_quat('camera1', sample_quat(CAM_ANGLE3))
fovy = sample(CAM_RFOVY)
self.cam_modder.set_fovy('camera1', fovy)
def _set_visible(self, prefix, range_top, visible):
"""Helper function to set visibility of several objects"""
if not visible:
if range_top == 0:
name = prefix
gid = self.model.geom_name2id(name)
self.model.geom_rgba[gid][-1] = 0.0
for i in range(range_top):
name = "{}{}".format(prefix, i)
gid = self.model.geom_name2id(name)
self.model.geom_rgba[gid][-1] = 0.0
else:
if range_top == 0:
name = prefix
gid = self.model.geom_name2id(name)
self.model.geom_rgba[gid][-1] = 1.0
for i in range(range_top):
name = "{}{}".format(prefix, i)
gid = self.model.geom_name2id(name)
self.model.geom_rgba[gid][-1] = 1.0
def mod_extra_distractors(self, visible=True):
"""mod rocks and tools on the side of the arena"""
# TODO: I might consider changing these to look like rocks instead of
# just random shapes. It just looks weird to me right now. Ok for now,
# but it seems a bit off.
N = 20
self._set_visible("side_obj", N, visible)
if not visible:
return
Z_JITTER = 0.05
OBJ_XRANGE = Range(0.01, 0.1)
OBJ_YRANGE = Range(0.01, 0.1)
OBJ_ZRANGE = Range(0.01, 0.1)
OBJ_SIZE_RANGE = Range3D(OBJ_XRANGE, OBJ_YRANGE, OBJ_ZRANGE)
floor_gid = self.model.geom_name2id("floor")
left_body_id = self.model.body_name2id("left_wall")
left_geom_id = self.model.geom_name2id("left_wall")
right_body_id = self.model.body_name2id("right_wall")
right_geom_id = self.model.geom_name2id("right_wall")
left_center = self.model.body_pos[left_body_id]
left_geo = self.model.geom_size[left_geom_id]
left_height = left_center[2] + left_geo[2]
left_xrange = Range(left_center[0] - left_geo[0],
left_center[0] + left_geo[0])
left_yrange = Range(left_center[1] - left_geo[1],
left_center[1] + left_geo[1])
left_zrange = 0.02 + Range(left_height - Z_JITTER,
left_height + Z_JITTER)
left_range = Range3D(left_xrange, left_yrange, left_zrange)
right_center = self.model.body_pos[right_body_id]
right_geo = self.model.geom_size[right_geom_id]
right_height = right_center[2] + right_geo[2]
right_xrange = Range(right_center[0] - right_geo[0],
right_center[0] + right_geo[0])
right_yrange = Range(right_center[1] - right_geo[1],
right_center[1] + right_geo[1])
right_zrange = 0.02 + Range(right_height - Z_JITTER,
right_height + Z_JITTER)
right_range = Range3D(right_xrange, right_yrange, right_zrange)
for i in range(N):
name = "side_obj{}".format(i)
obj_bid = self.model.body_name2id(name)
obj_gid = self.model.geom_name2id(name)
self.model.geom_quat[obj_gid] = random_quat()
self.model.geom_size[obj_gid] = sample_xyz(OBJ_SIZE_RANGE)
self.model.geom_type[obj_gid] = sample_geom_type()
# 50% chance of invisible
if sample([0, 1]) < 0.5:
self.model.geom_rgba[obj_gid][-1] = 0.0
else:
self.model.geom_rgba[obj_gid][-1] = 1.0
## 50% chance of same color as floor and rocks
if sample([0, 1]) < 0.5:
self.model.geom_matid[obj_gid] = self.model.geom_matid[
floor_gid]
else:
self.model.geom_matid[obj_gid] = self.start_matid[obj_gid]
# 10 always on the left, 10 always on the right
if i < 10:
self.model.body_pos[obj_bid] = sample_xyz(left_range)
else:
self.model.body_pos[obj_bid] = sample_xyz(right_range)
def mod_extra_judges(self, visible=True):
"""mod NASA judges around the perimeter of the arena"""
# TODO: might want to add regions on the sides of the arena, but these
# may be covered by the distractors already
N = 5
self._set_visible("judge", N, visible)
if not visible:
return
JUDGE_XRANGE = Range(0.1, 0.2)
JUDGE_YRANGE = Range(0.1, 0.2)
JUDGE_ZRANGE = Range(0.75, 1.0)
JUDGE_SIZE_RANGE = Range3D(JUDGE_XRANGE, JUDGE_YRANGE, JUDGE_ZRANGE)
digwall_bid = self.model.body_name2id("dig_wall")
digwall_gid = self.model.geom_name2id("dig_wall")
digwall_center = self.model.body_pos[digwall_bid]
digwall_geo = self.model.geom_size[digwall_gid]
digwall_xrange = Range(-1.0 + digwall_center[0] - digwall_geo[0],
1.0 + digwall_center[0] + digwall_geo[0])
digwall_yrange = Range(digwall_center[1] + 0.5,
digwall_center[1] + 1.5)
digwall_zrange = JUDGE_ZRANGE - 0.75
digwall_range = Range3D(digwall_xrange, digwall_yrange, digwall_zrange)
for i in range(N):
name = "judge{}".format(i)
judge_bid = self.model.body_name2id(name)
judge_gid = self.model.geom_name2id(name)
#self.model.geom_quat[judge_gid] = jitter_quat(self.start_geom_quat[judge_gid], 0.05)
self.model.geom_quat[judge_gid] = random_quat()
self.model.geom_size[judge_gid] = sample_xyz(JUDGE_SIZE_RANGE)
self.model.geom_type[judge_gid] = sample_geom_type()
if self.model.geom_type[judge_gid] == 3 or self.model.geom_type[
judge_gid] == 5:
self.model.geom_size[judge_gid][1] = self.model.geom_size[
judge_gid][2]
self.model.body_pos[judge_bid] = sample_xyz(digwall_range)
# 50% chance of invisible
self.model.geom_rgba[judge_gid][-1] = sample([0, 1]) < 0.5
def mod_extra_robot_parts(self, visible=True):
"""add distractor parts of robots in the lower area of the camera frame"""
N = 3
self._set_visible("robot_part", N, visible)
if not visible:
return
# Project difference into camera coordinate frame
cam_pos = self.model.cam_pos[0]
cam_quat = np.quaternion(*self.model.cam_quat[0])
lower_range = Range3D([0.0, 0.0], [-0.2, -0.3], [-0.2, -0.3])
lower_size = Range3D([0.2, 0.6], [0.01, 0.15], [0.01, 0.15])
lower_angle = Range3D([-85.0, -95.0], [-180, 180], [-85, -95])
upper_range = Range3D([-0.6, 0.6], [-0.05, 0.05], [-0.05, 0.05])
upper_size = Range3D([0.005, 0.05], [0.005, 0.05], [0.01, 0.3])
upper_angle = Range3D([-85.0, -95.0], [-180, 180], [-85, -95])
name = "robot_part0"
lower_bid = self.model.body_name2id(name)
lower_gid = self.model.geom_name2id(name)
lower_pos = cam_pos + quaternion.rotate_vectors(
cam_quat, sample_xyz(lower_range))
self.model.body_pos[lower_bid] = lower_pos
self.model.geom_size[lower_gid] = sample_xyz(lower_size)
self.model.geom_quat[lower_gid] = sample_quat(lower_angle)
self.model.geom_type[lower_gid] = sample_geom_type(reject=["capsule"])
if self.model.geom_type[lower_gid] == 5:
self.model.geom_size[lower_gid][0] = self.model.geom_size[
lower_gid][2]
for i in range(1, 10):
name = "robot_part{}".format(i)
upper_bid = self.model.body_name2id(name)
upper_gid = self.model.geom_name2id(name)
upper_pos = lower_pos + sample_xyz(upper_range)
self.model.body_pos[upper_bid] = upper_pos
self.model.geom_size[upper_gid] = sample_xyz(upper_size)
self.model.geom_type[upper_gid] = sample_geom_type()
# 50% of the time, choose random angle instead reasonable angle
if sample([0, 1]) < 0.5:
self.model.geom_quat[upper_gid] = sample_quat(upper_angle)
else:
self.model.geom_quat[upper_gid] = random_quat()
def mod_extra_arena_structure(self, visible=True):
"""add randomized structure of the arena in the background with
pillars and a crossbar"""
N = 16
self._set_visible("arena_structure", N, visible)
if not visible:
return
STARTY = DIGY + 3.0
ENDY = DIGY + 5.0
XBAR_SIZE = Range3D([10.0, 20.0], [0.05, 0.1], [0.2, 0.5])
XBAR_RANGE = Range3D([0.0, 0.0], [STARTY, ENDY], [3.0, 5.0])
STRUCTURE_SIZE = Range3D([0.05, 0.1], [0.05, 0.1], [3.0, 6.0])
STRUCTURE_RANGE = Range3D([np.nan, np.nan], [STARTY, ENDY], [0.0, 0.0])
x_range = np.linspace(ACX - 10.0, ACX + 10.0, 15)
# crossbar
name = "arena_structure{}".format(N - 1)
xbar_bid = self.model.body_name2id(name)
xbar_gid = self.model.geom_name2id(name)
self.model.geom_quat[xbar_gid] = jitter_quat(
self.start_geom_quat[xbar_gid], 0.01)
self.model.geom_size[xbar_gid] = sample_xyz(XBAR_SIZE)
self.model.body_pos[xbar_bid] = sample_xyz(XBAR_RANGE)
### 10% chance of invisible
##if sample([0,1]) < 0.1:
## self.model.geom_rgba[xbar_gid][-1] = 0.0
##else:
## self.model.geom_rgba[xbar_gid][-1] = 1.0
for i in range(N - 1):
name = "arena_structure{}".format(i)
arena_structure_bid = self.model.body_name2id(name)
arena_structure_gid = self.model.geom_name2id(name)
STRUCTURE_RANGE[0] = Range(x_range[i] - 0.1, x_range[i] + 0.1)
self.model.geom_quat[arena_structure_gid] = jitter_quat(
self.start_geom_quat[arena_structure_gid], 0.01)
self.model.geom_size[arena_structure_gid] = sample_xyz(
STRUCTURE_SIZE)
self.model.body_pos[arena_structure_bid] = sample_xyz(
STRUCTURE_RANGE)
self.model.geom_matid[arena_structure_gid] = self.model.geom_matid[
xbar_gid]
# 10% chance of invisible
if sample([0, 1]) < 0.1:
self.model.geom_rgba[arena_structure_gid][-1] = 0.0
else:
self.model.geom_rgba[arena_structure_gid][-1] = 1.0
def mod_extra_arena_background(self, visible=True, N=10):
""" """
self._set_visible("background", N, visible)
if not visible:
return
BACKGROUND_XRANGE = Range(0.1, 0.5)
BACKGROUND_YRANGE = Range(0.1, 0.5)
BACKGROUND_ZRANGE = Range(0.1, 1.0)
BACKGROUND_SIZE_RANGE = Range3D(BACKGROUND_XRANGE, BACKGROUND_YRANGE,
BACKGROUND_ZRANGE)
digwall_bid = self.model.body_name2id("dig_wall")
digwall_gid = self.model.geom_name2id("dig_wall")
c = digwall_center = self.model.body_pos[digwall_bid]
background_range = Range3D([c[0] - 4.0, c[0] + 4.0],
[c[1] + 8.0, c[1] + 12.0], [0.0, 5.0])
for i in range(N):
name = "background{}".format(i)
background_bid = self.model.body_name2id(name)
background_gid = self.model.geom_name2id(name)
self.model.geom_quat[background_gid] = random_quat()
self.model.geom_size[background_gid] = sample_xyz(
BACKGROUND_SIZE_RANGE)
self.model.geom_type[background_gid] = sample_geom_type()
self.model.body_pos[background_bid] = sample_xyz(background_range)
def mod_extra_lights(self, visible=True):
""""""
#visible = False
MODE_TARGETBODY = 3
STARTY = DIGY + 3.0
ENDY = DIGY + 5.0
LIGHT_RANGE = Range3D([ACX - 2.0, ACX + 2.0], [STARTY, ENDY],
[3.0, 5.0])
any_washed = False
for i in range(3):
name = "extra_light{}".format(i)
lid = self.model.light_name2id(name)
self.model.light_active[lid] = visible
if not visible:
return
self.model.light_pos[lid] = sample_xyz(LIGHT_RANGE)
self.model.light_dir[lid] = sample_light_dir()
self.model.light_directional[lid] = 0
### 3% chance of directional light, washing out colors
### (P(at least 1 will be triggered) ~= 0.1) = 1 - 3root(p)
washout = sample([0, 1]) < 0.03
any_washed = any_washed or washout
self.model.light_directional[lid] = washout
if any_washed:
self.mod_extra_light_discs(
visible=visible and sample([0, 1]) < 0.9)
else:
self.mod_extra_light_discs(visible=False)
def mod_extra_light_discs(self, visible=True):
""""""
N = 10
self._set_visible("light_disc", N, visible)
if not visible:
return
Z_JITTER = 0.05
DISC_XRANGE = Range(0.1, 4.0)
DISC_YRANGE = Range(0.1, 4.0)
DISC_ZRANGE = Range(0.1, 4.0)
DISC_SIZE_RANGE = Range3D(DISC_XRANGE, DISC_YRANGE, DISC_ZRANGE)
OUTR = 20.0
INR = 10.0
floor_bid = self.model.body_name2id("floor")
c = self.model.body_pos[floor_bid]
disc_xrange = Range(c[0] - OUTR, c[0] + OUTR)
disc_yrange = Range(c[1], c[1] + OUTR)
disc_zrange = Range(-5.0, 10.0)
disc_range = Range3D(disc_xrange, disc_yrange, disc_zrange)
for i in range(N):
name = "light_disc{}".format(i)
disc_bid = self.model.body_name2id(name)
disc_gid = self.model.geom_name2id(name)
disc_mid = self.model.geom_matid[disc_gid]
self.model.geom_quat[disc_gid] = random_quat()
self.model.geom_size[disc_gid] = sample_xyz(DISC_SIZE_RANGE)
self.model.geom_type[disc_gid] = sample_geom_type()
# keep trying to place the disc until it lands outside of blocking stuff
# (they will land in a u shape, because they are sampled with a
# constrain to not land in the middle)
while True:
xyz = sample_xyz(disc_range)
if ((xyz[0] > (c[0] - INR) and xyz[0] < (c[0] + INR))
and (xyz[1] > (c[1] - INR) and xyz[1] < (c[1] + INR))):
continue
else:
self.model.geom_pos[disc_gid] = xyz
break
# 50% chance of invisible
if sample([0, 1]) < 0.5:
self.model.geom_rgba[disc_gid][-1] = 0.0
else:
self.model.geom_rgba[disc_gid][-1] = 1.0
def mod_extras(self, visible=True):
"""
Randomize extra properties of the world such as the extra rocks on the side
of the arena wal
The motivation for these mods are that it seems likely that these distractor
objects, judges, etc, in the real images could degrade the performance
of the detector.
Artifacts:
- Rocks and tools on edges of bin
- NASA judges around the perimeter
- Parts of the robot in the frame
- Background arena structure and crowd
- Bright light around edges of arena
"""
self.mod_extra_distractors(visible)
self.mod_extra_robot_parts(visible)
self.mod_extra_lights(visible)
# 10% of the time, hide the other distractor pieces
self.mod_extra_judges(visible=visible and sample([0, 1]) < 0.9)
self.mod_extra_arena_structure(
visible=visible and sample([0, 1]) < 0.9)
self.mod_extra_arena_background(
visible=visible and sample([0, 1]) < 0.9)
def mod_walls(self):
"""
Randomize the x, y, and orientation of the walls slights.
Also drastically randomize the height of the walls. In many cases they won't
be seen at all. This will allow the model to generalize to scenarios without
walls, or where the walls and geometry is slightly different than the sim
model
"""
wall_names = ["left_wall", "right_wall", "bin_wall", "dig_wall"]
for name in wall_names:
geom_id = self.model.geom_name2id(name)
body_id = self.model.body_name2id(name)
jitter_x = Range(-0.2, 0.2)
jitter_y = Range(-0.2, 0.2)
jitter_z = Range(-0.75, 0.0)
jitter3D = Range3D(jitter_x, jitter_y, jitter_z)
self.model.body_pos[
body_id] = self.start_body_pos[body_id] + sample_xyz(jitter3D)
self.model.body_quat[body_id] = jitter_quat(
self.start_body_quat[body_id], 0.005)
if sample([0, 1]) < 0.05:
self.model.body_pos[body_id][2] = -2.0
def mod_dirt(self):
"""Randomize position and rotation of dirt"""
# TODO: 50% chance to flip the dirt pile upsidedown, then we will have
# to account for this in calculations
# dirt stuff
DIRT_RX = Range(0.0, 0.3)
DIRT_RY = Range(0.0, 0.3)
DIRT_RZ = Range(-0.05, 0.03)
DIRT_RANGE3D = Range3D(DIRT_RX, DIRT_RY, DIRT_RZ)
DIRT_RYAW = Range(-180, 180)
DIRT_RPITCH = Range(-90.5, -89.5)
DIRT_RROLL = Range(-0.5, 0.5)
DIRT_ANGLE3 = Range3D(DIRT_RYAW, DIRT_RPITCH, DIRT_RROLL)
dirt_bid = self.model.body_name2id("dirt")
dirt_gid = self.model.geom_name2id("dirt")
dirt_mid = self.model.geom_dataid[dirt_gid]
# randomize position and yaw of dirt
self.model.body_pos[dirt_bid] = self.start_body_pos[
dirt_bid] + sample_xyz(DIRT_RANGE3D)
self.model.geom_quat[dirt_gid] = sample_quat(DIRT_ANGLE3)
vert_adr = self.model.mesh_vertadr[dirt_mid]
vert_num = self.model.mesh_vertnum[dirt_mid]
mesh_verts = self.model.mesh_vert[vert_adr:vert_adr + vert_num]
rot_quat = self.model.geom_quat[dirt_gid]
rots = quaternion.rotate_vectors(
np.quaternion(*rot_quat).normalized(), mesh_verts)
mesh_abs_pos = self.floor_offset + self.model.body_pos[dirt_bid] + rots
#xy_indexes = mesh_abs_pos[:, 0:2]
#z_heights = mesh_abs_pos[:, 2]
# Return a function that the user can call to get the approximate
# height of an xy location
def local_mean_height(xy):
"""
Take an xy coordinate, and the approximate z height of the mesh at that
location. It works decently.
Uses a weighted average mean of all points within a threshold of xy.
"""
# grab all mesh points within threshold euclidean distance of xy
gt0_xyz = mesh_abs_pos[mesh_abs_pos[:, 2] > 0.01]
eudists = np.sum(np.square(gt0_xyz[:, 0:2] - xy), axis=1)
indices = eudists < 0.1
close_xyz = gt0_xyz[indices]
# if there are any nearby points above 0
if np.count_nonzero(close_xyz[:, 2]) > 0:
# weights for weighted sum. closer points to xy have higher weight
weights = 1 / (eudists[indices])
weights = np.expand_dims(weights / np.sum(weights), axis=1)
pos = np.sum(close_xyz * weights, axis=0)
# show an "o" and a marker where the height is
if self.visualize:
self.viewer.add_marker(pos=pos,
label="o",
size=np.array([0.01, 0.01, 0.01]),
rgba=np.array([0.0, 1.0, 0.0, 1.0]))
# approximate z height of ground
return pos[2]
else:
return 0
def always_zero(xy):
return 0
dirt_height_xy = local_mean_height
#dirt_height_xy = always_zero
return dirt_height_xy
def mod_rocks(self):
"""
Randomize the rocks so that the model will generalize to competition rocks
This randomizations currently being done are:
- Positions (within guesses of competition regions)
- Orientations
- Shuffling the 3 rock meshes so that they can be on the left, middle, or right
- Generating new random rock meshes every n runs (with Blender)
"""
# Rock placement range parameters
ROCK_LANEX = 0.4 # width parameters of x range
OUTER_EXTRA = 0.5 # how much farther rocks should go out on the right and left lanes
ROCK_BUFFX = 0.2 # distacne between rock lanes
# How far into the obstacle zone the rocks should start.
ROCK_START_OFFSET = 0.2
MID_START_OFFSET = 0.4 # bit more for middle rock
ROCK_RY = Range(OBS_SY + ROCK_START_OFFSET, OBS_ENDY)
MID_RY = Range(OBS_SY + MID_START_OFFSET, OBS_ENDY)
ROCK_RZ = Range(AFZ - 0.02, AFZ + 0.2)
# Position dependent ranges
LEFT_RX = Range(-3 * ROCK_LANEX - OUTER_EXTRA,
-ROCK_LANEX - ROCK_BUFFX)
MID_RX = Range(-ROCK_LANEX, ROCK_LANEX)
RIGHT_RX = Range(ROCK_BUFFX + ROCK_LANEX, 3 * ROCK_LANEX + OUTER_EXTRA)
# Form full 3D sample range
LEFT_ROCK_RANGE = Range3D(LEFT_RX, ROCK_RY, ROCK_RZ)
MID_ROCK_RANGE = Range3D(MID_RX, MID_RY, ROCK_RZ)
RIGHT_ROCK_RANGE = Range3D(RIGHT_RX, ROCK_RY, ROCK_RZ)
ROCK_RANGES = [LEFT_ROCK_RANGE, MID_ROCK_RANGE, RIGHT_ROCK_RANGE]
# actual mods
rock_body_ids = {}
rock_geom_ids = {}
rock_mesh_ids = {}
max_height_idxs = {}
rot_cache = {}
#max_height_xys = {}
dirt_height_xy = self.mod_dirt()
for name in self.model.geom_names:
if name[:4] != "rock":
continue
geom_id = self.model.geom_name2id(name)
body_id = self.model.body_name2id(name)
mesh_id = self.model.geom_dataid[geom_id]
rock_geom_ids[name] = geom_id
rock_body_ids[name] = body_id
rock_mesh_ids[name] = mesh_id
# Rotate the rock and get the z value of the highest point in the
# rotated rock mesh
rot_quat = random_quat()
vert_adr = self.model.mesh_vertadr[mesh_id]
vert_num = self.model.mesh_vertnum[mesh_id]
mesh_verts = self.model.mesh_vert[vert_adr:vert_adr + vert_num]
rots = quaternion.rotate_vectors(
np.quaternion(*rot_quat).normalized(), mesh_verts)
self.model.geom_quat[geom_id] = rot_quat
max_height_idx = np.argmax(rots[:, 2])
max_height_idxs[name] = max_height_idx
rot_cache[name] = rots
# Shuffle the positions of the rocks (l or m or r)
shuffle_names = list(rock_body_ids.keys())
random.shuffle(shuffle_names)
self.rock_mod_cache = []
for i in range(len(shuffle_names)):
name = shuffle_names[i]
rots = rot_cache[name]
self.model.body_pos[rock_body_ids[name]] = np.array(
sample_xyz(ROCK_RANGES[i]))
max_height_idx = max_height_idxs[name]
xyz_for_max_z = rots[max_height_idx]
# xyz coords in global frame
global_xyz = self.floor_offset + xyz_for_max_z + self.model.body_pos[
rock_body_ids[name]]
gxy = global_xyz[0:2]
max_height = global_xyz[2]
if self.visualize:
self.viewer.add_marker(pos=global_xyz,
label="m",
size=np.array([0.01, 0.01, 0.01]),
rgba=np.array([0.0, 0.0, 1.0, 1.0]))
dirt_z = dirt_height_xy(gxy)
#dirt_z = 0
#print(name, dirt_z)
z_height = max_height - dirt_z
self.rock_mod_cache.append((name, z_height))
def get_ground_truth(self):
"""
self.rock_mod_cache set from self.mod_rocks
Return 1d numpy array of 9 elements for positions of all 3 rocks including:
- rock x dist from cam
- rock y dist from cam
- rock z height from arena floor
"""
cam_pos = self.model.cam_pos[0]
#line_pos = self.floor_offset + np.array([0.0, 0.75, 0.0])
#self.viewer.add_marker(pos=line_pos)
##r0_pos = self.floor_offset + self.model.body_pos[self.model.body_name2id('rock0')]
##r1_pos = self.floor_offset + self.model.body_pos[self.model.body_name2id('rock1')]
##r2_pos = self.floor_offset + self.model.body_pos[self.model.body_name2id('rock2')]
##r0_diff = r0_pos - cam_pos
##r1_diff = r1_pos - cam_pos
##r2_diff = r2_pos - cam_pos
ground_truth = np.zeros(9, dtype=np.float32)
for i, slot in enumerate(self.rock_mod_cache):
name = slot[0]
z_height = slot[1]
pos = self.floor_offset + self.model.body_pos[
self.model.body_name2id(name)]
diff = pos - cam_pos
# Project difference into camera coordinate frame
cam_angle = quaternion.as_euler_angles(
np.quaternion(*self.model.cam_quat[0]))[0]
cam_angle += np.pi / 2
in_cam_frame = np.zeros_like(diff)
x = diff[1]
y = -diff[0]
in_cam_frame[0] = x * np.cos(cam_angle) + y * np.sin(cam_angle)
in_cam_frame[1] = -x * np.sin(cam_angle) + y * np.cos(cam_angle)
in_cam_frame[2] = z_height
# simple check that change of frame is mathematically valid
assert (np.isclose(np.sum(np.square(diff[:2])),
np.sum(np.square(in_cam_frame[:2]))))
# swap positions to match ROS standard coordinates
ground_truth[3 * i + 0] = in_cam_frame[0]
ground_truth[3 * i + 1] = in_cam_frame[1]
ground_truth[3 * i + 2] = in_cam_frame[2]
text = "{0} x: {1:.2f} y: {2:.2f} height:{3:.2f}".format(
name, ground_truth[3 * i + 0], ground_truth[3 * i + 1],
z_height)
##text = "x: {0:.2f} y: {1:.2f} height:{2:.2f}".format(ground_truth[3*i+0], ground_truth[3*i+1], z_height)
##text = "height:{0:.2f}".format(z_height)
if self.visualize:
self.viewer.add_marker(pos=pos, label=text, rgba=np.zeros(4))
return ground_truth
def __init__(
self,
reward_info=None,
frame_skip=50,
pos_action_scale=4. / 100,
randomize_goals=True,
puck_goal_low=(-0.1, 0.5),
puck_goal_high=(0.1, 0.7),
hand_goal_low=(-0.1, 0.5),
hand_goal_high=(0.1, 0.7),
mocap_low=(-0.1, 0.5, 0.0),
mocap_high=(0.1, 0.7, 0.5),
# unused
init_block_low=(-0.05, 0.55),
init_block_high=(0.05, 0.65),
fixed_puck_goal=(0.05, 0.6),
fixed_hand_goal=(-0.05, 0.6),
# multi-object
num_objects=3,
fixed_colors=True,
seed=1,
filename='sawyer_multiobj.xml',
object_mass=1,
# object_meshes=['Bowl', 'GlassBowl', 'LotusBowl01', 'ElephantBowl', 'RuggedBowl'],
object_meshes=None,
obj_classname=None,
block_height=0.02,
block_width=0.02,
cylinder_radius=0.035,
finger_sensors=False,
maxlen=0.07,
minlen=0.01,
preload_obj_dict=None,
reset_to_initial_position=True,
object_low=(-np.inf, -np.inf, -np.inf),
object_high=(np.inf, np.inf, np.inf),
action_repeat=1,
fixed_start=False,
fixed_start_pos=(0, 0.6),
goal_moves_one_object=True,
num_scene_objects=[
0, 1, 2
], # list of number of objects that can appear per scene
object_height=0.02,
use_textures=False,
init_camera=None,
sliding_joints=False,
):
if seed:
np.random.seed(seed)
self.env_seed = seed
self.quick_init(locals())
self.reward_info = reward_info
self.randomize_goals = randomize_goals
self._pos_action_scale = pos_action_scale
self.reset_to_initial_position = reset_to_initial_position
self.init_block_low = np.array(init_block_low)
self.init_block_high = np.array(init_block_high)
self.puck_goal_low = np.array(puck_goal_low)
self.puck_goal_high = np.array(puck_goal_high)
self.hand_goal_low = np.array(hand_goal_low)
self.hand_goal_high = np.array(hand_goal_high)
self.fixed_puck_goal = np.array(fixed_puck_goal)
self.fixed_hand_goal = np.array(fixed_hand_goal)
self.mocap_low = np.array(mocap_low)
self.mocap_high = np.array(mocap_high)
self.object_low = np.array(object_low)
self.object_high = np.array(object_high)
self.action_repeat = action_repeat
self.fixed_colors = fixed_colors
self.goal_moves_one_object = goal_moves_one_object
self.num_objects = num_objects
self.num_scene_objects = num_scene_objects
self.object_height = object_height
self.fixed_start = fixed_start
self.fixed_start_pos = np.array(fixed_start_pos)
self.maxlen = maxlen
self.use_textures = use_textures
self.sliding_joints = sliding_joints
self.cur_objects = np.array([0, 1, 2])
self.preload_obj_dict = preload_obj_dict
self.num_cur_objects = 0
# Generate XML
base_filename = asset_base_path + filename
friction_params = (0.1, 0.1, 0.02)
self.obj_stat_prop = create_object_xml(
base_filename, num_objects, object_mass, friction_params,
object_meshes, finger_sensors, maxlen, minlen, preload_obj_dict,
obj_classname, block_height, block_width, cylinder_radius,
use_textures, sliding_joints)
gen_xml = create_root_xml(base_filename)
MujocoEnv.__init__(self, gen_xml, frame_skip=frame_skip)
clean_xml(gen_xml)
if self.use_textures:
self.modder = TextureModder(self.sim)
self.state_goal = self.sample_goal_for_rollout()
# MultitaskEnv.__init__(self, distance_metric_order=2)
# MujocoEnv.__init__(self, gen_xml, frame_skip=frame_skip)
self.action_space = Box(
np.array([-1, -1]),
np.array([1, 1]),
)
self.num_objects = num_objects
low = (self.num_scene_objects[0] + 1) * [-0.2, 0.5]
high = (self.num_scene_objects[0] + 1) * [0.2, 0.7]
self.obs_box = Box(
np.array(low),
np.array(high),
)
goal_low = np.array(
low) # np.concatenate((self.hand_goal_low, self.puck_goal_low))
goal_high = np.array(
high) # np.concatenate((self.hand_goal_high, self.puck_goal_high))
self.goal_box = Box(
goal_low,
goal_high,
)
self.total_objects = self.num_objects + 1
self.objects_box = Box(
np.zeros((self.total_objects, )),
np.ones((self.total_objects, )),
)
self.observation_space = Dict([
('observation', self.obs_box),
('state_observation', self.obs_box),
('desired_goal', self.goal_box),
('state_desired_goal', self.goal_box),
('achieved_goal', self.goal_box),
('state_achieved_goal', self.goal_box),
('objects', self.objects_box),
])
# hack for state-based experiments for other envs
# self.observation_space = Box(
# np.array([-0.2, 0.5, -0.2, 0.5, -0.2, 0.5]),
# np.array([0.2, 0.7, 0.2, 0.7, 0.2, 0.7]),
# )
# self.goal_space = Box(
# np.array([-0.2, 0.5, -0.2, 0.5, -0.2, 0.5]),
# np.array([0.2, 0.7, 0.2, 0.7, 0.2, 0.7]),
# )
self.set_initial_object_positions()
if use_textures:
super().initialize_camera(init_camera)
self.initialized_camera = init_camera
self.reset()
self.reset_mocap_welds()
class SawyerMultiobjectEnv(MujocoEnv, Serializable, MultitaskEnv):
INIT_HAND_POS = np.array([0, 0.4, 0.02])
def __init__(
self,
reward_info=None,
frame_skip=50,
pos_action_scale=4. / 100,
randomize_goals=True,
puck_goal_low=(-0.1, 0.5),
puck_goal_high=(0.1, 0.7),
hand_goal_low=(-0.1, 0.5),
hand_goal_high=(0.1, 0.7),
mocap_low=(-0.1, 0.5, 0.0),
mocap_high=(0.1, 0.7, 0.5),
# unused
init_block_low=(-0.05, 0.55),
init_block_high=(0.05, 0.65),
fixed_puck_goal=(0.05, 0.6),
fixed_hand_goal=(-0.05, 0.6),
# multi-object
num_objects=3,
fixed_colors=True,
seed=1,
filename='sawyer_multiobj.xml',
object_mass=1,
# object_meshes=['Bowl', 'GlassBowl', 'LotusBowl01', 'ElephantBowl', 'RuggedBowl'],
object_meshes=None,
obj_classname=None,
block_height=0.02,
block_width=0.02,
cylinder_radius=0.035,
finger_sensors=False,
maxlen=0.07,
minlen=0.01,
preload_obj_dict=None,
reset_to_initial_position=True,
object_low=(-np.inf, -np.inf, -np.inf),
object_high=(np.inf, np.inf, np.inf),
action_repeat=1,
fixed_start=False,
fixed_start_pos=(0, 0.6),
goal_moves_one_object=True,
num_scene_objects=[
0, 1, 2
], # list of number of objects that can appear per scene
object_height=0.02,
use_textures=False,
init_camera=None,
sliding_joints=False,
):
if seed:
np.random.seed(seed)
self.env_seed = seed
self.quick_init(locals())
self.reward_info = reward_info
self.randomize_goals = randomize_goals
self._pos_action_scale = pos_action_scale
self.reset_to_initial_position = reset_to_initial_position
self.init_block_low = np.array(init_block_low)
self.init_block_high = np.array(init_block_high)
self.puck_goal_low = np.array(puck_goal_low)
self.puck_goal_high = np.array(puck_goal_high)
self.hand_goal_low = np.array(hand_goal_low)
self.hand_goal_high = np.array(hand_goal_high)
self.fixed_puck_goal = np.array(fixed_puck_goal)
self.fixed_hand_goal = np.array(fixed_hand_goal)
self.mocap_low = np.array(mocap_low)
self.mocap_high = np.array(mocap_high)
self.object_low = np.array(object_low)
self.object_high = np.array(object_high)
self.action_repeat = action_repeat
self.fixed_colors = fixed_colors
self.goal_moves_one_object = goal_moves_one_object
self.num_objects = num_objects
self.num_scene_objects = num_scene_objects
self.object_height = object_height
self.fixed_start = fixed_start
self.fixed_start_pos = np.array(fixed_start_pos)
self.maxlen = maxlen
self.use_textures = use_textures
self.sliding_joints = sliding_joints
self.cur_objects = np.array([0, 1, 2])
self.preload_obj_dict = preload_obj_dict
self.num_cur_objects = 0
# Generate XML
base_filename = asset_base_path + filename
friction_params = (0.1, 0.1, 0.02)
self.obj_stat_prop = create_object_xml(
base_filename, num_objects, object_mass, friction_params,
object_meshes, finger_sensors, maxlen, minlen, preload_obj_dict,
obj_classname, block_height, block_width, cylinder_radius,
use_textures, sliding_joints)
gen_xml = create_root_xml(base_filename)
MujocoEnv.__init__(self, gen_xml, frame_skip=frame_skip)
clean_xml(gen_xml)
if self.use_textures:
self.modder = TextureModder(self.sim)
self.state_goal = self.sample_goal_for_rollout()
# MultitaskEnv.__init__(self, distance_metric_order=2)
# MujocoEnv.__init__(self, gen_xml, frame_skip=frame_skip)
self.action_space = Box(
np.array([-1, -1]),
np.array([1, 1]),
)
self.num_objects = num_objects
low = (self.num_scene_objects[0] + 1) * [-0.2, 0.5]
high = (self.num_scene_objects[0] + 1) * [0.2, 0.7]
self.obs_box = Box(
np.array(low),
np.array(high),
)
goal_low = np.array(
low) # np.concatenate((self.hand_goal_low, self.puck_goal_low))
goal_high = np.array(
high) # np.concatenate((self.hand_goal_high, self.puck_goal_high))
self.goal_box = Box(
goal_low,
goal_high,
)
self.total_objects = self.num_objects + 1
self.objects_box = Box(
np.zeros((self.total_objects, )),
np.ones((self.total_objects, )),
)
self.observation_space = Dict([
('observation', self.obs_box),
('state_observation', self.obs_box),
('desired_goal', self.goal_box),
('state_desired_goal', self.goal_box),
('achieved_goal', self.goal_box),
('state_achieved_goal', self.goal_box),
('objects', self.objects_box),
])
# hack for state-based experiments for other envs
# self.observation_space = Box(
# np.array([-0.2, 0.5, -0.2, 0.5, -0.2, 0.5]),
# np.array([0.2, 0.7, 0.2, 0.7, 0.2, 0.7]),
# )
# self.goal_space = Box(
# np.array([-0.2, 0.5, -0.2, 0.5, -0.2, 0.5]),
# np.array([0.2, 0.7, 0.2, 0.7, 0.2, 0.7]),
# )
self.set_initial_object_positions()
if use_textures:
super().initialize_camera(init_camera)
self.initialized_camera = init_camera
self.reset()
self.reset_mocap_welds()
def initialize_camera(self, init_fctn):
if self.use_textures:
# do nothing, because the camera was already initialized
assert init_fctn == self.initialized_camera, "cameras do not match"
else:
super().initialize_camera(init_fctn)
@property
def model_name(self):
return get_asset_full_path(
'sawyer_xyz/sawyer_push_and_reach_mocap_goal_hidden.xml')
def viewer_setup(self):
self.viewer.cam.trackbodyid = 0
self.viewer.cam.distance = 1.0
# robot view
# rotation_angle = 90
# cam_dist = 1
# cam_pos = np.array([0, 0.5, 0.2, cam_dist, -45, rotation_angle])
# 3rd person view
cam_dist = 0.3
rotation_angle = 270
cam_pos = np.array([0, 1.0, 0.5, cam_dist, -45, rotation_angle])
# top down view
# cam_dist = 0.2
# rotation_angle = 0
# cam_pos = np.array([0, 0, 1.5, cam_dist, -90, rotation_angle])
for i in range(3):
self.viewer.cam.lookat[i] = cam_pos[i]
self.viewer.cam.distance = cam_pos[3]
self.viewer.cam.elevation = cam_pos[4]
self.viewer.cam.azimuth = cam_pos[5]
self.viewer.cam.trackbodyid = -1
def step(self, a):
a = np.clip(a, -1, 1)
mocap_delta_z = 0.06 - self.data.mocap_pos[0, 2]
new_mocap_action = np.hstack((a, np.array([mocap_delta_z])))
u = np.zeros(7)
try:
for _ in range(self.action_repeat):
self.mocap_set_action(new_mocap_action[:3] *
self._pos_action_scale)
self.do_simulation(u, self.frame_skip)
except mujoco_py.builder.MujocoException:
pass
# self.render()
qpos = self.data.qpos.flat.copy()
qvel = self.data.qvel.flat.copy()
for i in range(self.num_objects):
#if i in self.cur_objects:
x = 7 + i * 7
y = 10 + i * 7
qpos[x:y] = np.clip(qpos[x:y], self.object_low, self.object_high)
self.set_state(qpos, qvel)
endeff_pos = self.get_endeff_pos()
hand_distance = np.linalg.norm(self.get_hand_goal_pos() - endeff_pos)
object_distances = {}
touch_distances = {}
# import ipdb; ipdb.set_trace()
# for i in range(self.num_objects):
# if i in self.cur_objects:
i = 0
object_name = "object%d_distance" % i
object_distance = np.linalg.norm(
self.get_object_goal_pos(i) - self.get_object_pos(i))
object_distances[object_name] = object_distance
touch_name = "touch%d_distance" % i
touch_distance = np.linalg.norm(endeff_pos - self.get_object_pos(i))
touch_distances[touch_name] = touch_distance
objects = {}
# for i in range(self.num_objects):
distances = []
cur_object_list = self.cur_objects.tolist()
for i in self.cur_objects:
j = cur_object_list.index(i)
object_goal = self.get_object_goal_pos(j)
object_pos = self.get_object_pos(i)
object_distance = np.linalg.norm(object_pos - object_goal)
distances.append(object_distance)
object_distances["current_object_distance"] = np.mean(distances)
b = np.zeros((self.num_objects + 1))
b[0] = 1 # the arm
for i in self.cur_objects:
b[i + 1] = 1
for i in range(self.num_objects):
objects["object%d" % i] = b[i + 1]
info = dict(
hand_distance=hand_distance,
success=float(
hand_distance + sum(object_distances.values()) < 0.06),
**object_distances,
**touch_distances,
**objects,
objects_present=b,
)
obs = self._get_obs()
# reward = self.compute_reward(obs, u, obs, self._goal_xyxy)
reward = self.compute_rewards(a, obs, info)
done = False
return obs, reward, done, info
def mocap_set_action(self, action):
pos_delta = action[None]
new_mocap_pos = self.data.mocap_pos + pos_delta
new_mocap_pos[0, :] = np.clip(new_mocap_pos[0, :], self.mocap_low,
self.mocap_high)
self.data.set_mocap_pos('mocap', new_mocap_pos)
self.data.set_mocap_quat('mocap', np.array([1, 0, 1, 0]))
def _get_obs(self):
e = self.get_endeff_pos()[:2]
bs = []
for i in range(self.num_objects):
if i in self.cur_objects:
b = self.get_object_pos(i)[:2]
bs.append(b)
b = np.concatenate(bs)
x = np.concatenate((e, b))
g = self.state_goal
o = np.zeros((self.total_objects))
o[0] = 1 # the hand
for i in self.cur_objects:
o[i + 1] = 1
new_obs = dict(
observation=x,
state_observation=x,
desired_goal=g,
state_desired_goal=g,
achieved_goal=x,
state_achieved_goal=x,
objects=o,
)
return new_obs
def get_object_pos(self, id):
mujoco_id = self.model.body_names.index('object' + str(id))
return self.data.body_xpos[mujoco_id].copy()[:2]
def get_endeff_pos(self):
return self.data.body_xpos[self.endeff_id].copy()[:2]
def get_hand_goal_pos(self):
return self.state_goal[:2]
def get_object_goal_pos(self, i):
x = 2 + 2 * i
y = 4 + 2 * i
return self.state_goal[x:y]
@property
def endeff_id(self):
return self.model.body_names.index('leftclaw')
def set_object_xy(self, i, pos):
qpos = self.data.qpos.flat.copy()
qvel = self.data.qvel.flat.copy()
x = 7 + i * 7
y = 10 + i * 7
z = 14 + i * 7
qpos[x:y] = np.hstack((pos.copy(), np.array([self.object_height])))
qpos[y:z] = np.array([1, 0, 0, 0])
x = 7 + i * 6
y = 13 + i * 6
qvel[x:y] = 0
self.set_state(qpos, qvel)
def set_object_xys(self, positions):
qpos = self.data.qpos.flat.copy()
qvel = self.data.qvel.flat.copy()
for i, pos in enumerate(positions):
x = 7 + i * 7
y = 10 + i * 7
z = 14 + i * 7
qpos[x:y] = np.hstack((pos.copy(), np.array([self.object_height])))
qpos[y:z] = np.array([1, 0, 0, 0])
x = 7 + i * 6
y = 13 + i * 6
qvel[x:y] = 0
self.set_state(qpos, qvel)
def reset_mocap_welds(self):
"""Resets the mocap welds that we use for actuation."""
sim = self.sim
if sim.model.nmocap > 0 and sim.model.eq_data is not None:
for i in range(sim.model.eq_data.shape[0]):
if sim.model.eq_type[i] == mujoco_py.const.EQ_WELD:
sim.model.eq_data[i, :] = np.array(
[0., 0., 0., 1., 0., 0., 0.])
sim.forward()
def reset_mocap2body_xpos(self):
# move mocap to weld joint
self.data.set_mocap_pos(
'mocap',
np.array([self.data.body_xpos[self.endeff_id]]),
)
self.data.set_mocap_quat(
'mocap',
np.array([self.data.body_xquat[self.endeff_id]]),
)
def set_initial_object_positions(self):
n_o = np.random.choice(self.num_scene_objects)
n_o = len(self.num_scene_objects)
self.num_cur_objects = n_o
#self.cur_objects = np.random.choice(self.num_objects, n_o, replace=False)
#self.cur_objects = [0, 1, 2]
while True:
pos = [
self.INIT_HAND_POS[:2],
]
for i in range(n_o):
if self.fixed_start:
r = self.fixed_start_pos
else:
r = np.random.uniform(self.puck_goal_low,
self.puck_goal_high)
pos.append(r)
touching = []
for i in range(n_o + 1):
for j in range(i):
t = np.linalg.norm(pos[i] - pos[j]) <= self.maxlen
touching.append(t)
if not any(touching):
break
pos.reverse()
positions = []
for i in range(self.num_objects):
z = 10 + 10 * i
positions.append(np.array([z, z]))
for i in range(n_o):
j = self.cur_objects[i]
positions[j] = pos[i]
self.set_object_xys(positions)
def reset(self):
if self.use_textures:
for i in range(self.num_objects):
self.modder.rand_rgb('object%d' % i)
velocities = self.data.qvel.copy()
angles = self.data.qpos.copy()
angles[:7] = np.array(self.init_angles[:7]) # just change robot joints
self.set_state(angles.flatten(), velocities.flatten())
for _ in range(10):
self.data.set_mocap_pos('mocap', self.INIT_HAND_POS)
self.data.set_mocap_quat('mocap', np.array([1, 0, 1, 0]))
# set_state resets the goal xy, so we need to explicit set it again
# if self.reset_to_initial_position:
# self.set_initial_object_positions()
self.set_initial_object_positions()
self.state_goal = self.sample_goal_for_rollout()
self.reset_mocap_welds()
# import ipdb; ipdb.set_trace()
# for i in range(self.num_objects):
# obj_id = self.model.body_names.index('object0')
# xpos = self.data.body_xpos[obj_id]
# xquat = self.data.body_xquat[obj_id]
# self.data.set_joint_qpos(xpos)
#self.set_initial_object_positions()
return self._get_obs()
def compute_rewards(self, action, obs, info=None):
# objects_present = info['objects_present'].reshape(-1, self.num_objects + 1, 1)
ob_p = obs['state_achieved_goal'].reshape(-1, self.num_cur_objects + 1,
2)
goal = obs['state_desired_goal'].reshape(-1, self.num_cur_objects + 1,
2)
# th = objects_present*ob_p != 0
# ob = ob_p[:, th[0][:, 0]]
distances = np.linalg.norm(ob_p - goal, axis=2)[:, 1:]
return -distances
# def compute_reward(self, action, obs, info=None):
# r = -np.linalg.norm(obs['state_achieved_goal'] - obs['state_desired_goal'])
# return r
def compute_her_reward_np(self, ob, action, next_ob, goal, env_info=None):
return self.compute_reward(ob,
action,
next_ob,
goal,
env_info=env_info)
@property
def init_angles(self):
return [
1.78026069e+00,
-6.84415781e-01,
-1.54549231e-01,
2.30672090e+00,
1.93111471e+00,
1.27854012e-01,
1.49353907e+00,
1.80196716e-03,
7.40415706e-01,
2.09895360e-02,
9.99999990e-01,
3.05766105e-05,
-3.78462492e-06,
1.38684523e-04,
-3.62518873e-02,
6.13435141e-01,
2.09686080e-02,
7.07106781e-01,
1.48979724e-14,
7.07106781e-01,
-1.48999170e-14,
0,
0.6,
0.02,
1,
0,
1,
0,
]
def log_diagnostics(self, paths, logger=None, prefix=""):
if logger is None:
return
statistics = OrderedDict()
# for stat_name in [
# 'hand_distance',
# 'success',
# ] + [
# "object%d_distance" % i for i in range(self.num_objects)
# ] + [
# "touch%d_distance" % i for i in range(self.num_objects)
# ] + [
# "object%d" % i for i in range(self.num_objects)
# ]:
# stat_name = stat_name
# stat = get_stat_in_paths(paths, 'env_infos', stat_name)
# statistics.update(create_stats_ordered_dict(
# '%s %s' % (prefix, stat_name),
# stat,
# always_show_all_stats=True,
# ))
# statistics.update(create_stats_ordered_dict(
# 'Final %s %s' % (prefix, stat_name),
# [s[-1] for s in stat],
# always_show_all_stats=True,
# ))
for key, value in statistics.items():
logger.record_tabular(key, value)
"""
Multitask functions
"""
@property
def goal_dim(self) -> int:
return 4
def sample_goals(self, batch_size):
goals = np.random.uniform(
self.goal_box.low,
self.goal_box.high,
size=(batch_size, self.goal_box.low.size),
)
return {
'desired_goal': goals,
'state_desired_goal': goals,
}
def sample_goal_for_rollout(self):
n = len(self.cur_objects)
if self.randomize_goals:
if self.goal_moves_one_object:
hand = np.random.uniform(self.hand_goal_low,
self.hand_goal_high)
bs = []
for i in range(self.num_objects):
if i in self.cur_objects:
b = self.get_object_pos(i)[:2]
bs.append(b)
if n:
r = np.random.choice(self.cur_objects) # object to move
pos = bs + [
self.INIT_HAND_POS[:2],
]
while True:
bs[r] = np.random.uniform(self.puck_goal_low,
self.puck_goal_high)
touching = []
for i in range(n + 1):
if i != r:
t = np.linalg.norm(pos[i] -
bs[r]) <= self.maxlen
touching.append(t)
if not any(touching):
break
puck = np.concatenate(bs)
else:
hand = np.random.uniform(self.hand_goal_low,
self.hand_goal_high)
puck = np.concatenate([
np.random.uniform(self.puck_goal_low, self.puck_goal_high)
for i in range(n)
])
else:
hand = self.fixed_hand_goal.copy()
puck = self.fixed_puck_goal.copy()
g = np.hstack((hand, puck))
return g
# OLD SET GOAL
# def set_goal(self, goal):
# MultitaskEnv.set_goal(self, goal)
# self.set_goal_xyxy(goal)
# # hack for VAE
# self.set_to_goal(goal)
def get_goal(self):
return {
'desired_goal': self.state_goal,
'state_desired_goal': self.state_goal,
}
def set_goal(self, goal):
self.state_goal = goal['state_desired_goal']
def set_to_goal(self, goal):
state_goal = goal['state_desired_goal']
self.set_hand_xy(state_goal[:2])
# disappear all the objects
for i in range(self.num_objects):
z = 10 + 10 * i
self.set_object_xy(i, np.array([z, z]))
# set the goal positions of only the current objects
for i, j in enumerate(self.cur_objects):
x = 2 + 2 * i
y = 4 + 2 * i
self.set_object_xy(j, state_goal[x:y])
def convert_obs_to_goals(self, obs):
return obs
def set_hand_xy(self, xy):
for _ in range(10):
self.data.set_mocap_pos('mocap', np.array([xy[0], xy[1], 0.02]))
self.data.set_mocap_quat('mocap', np.array([1, 0, 1, 0]))
u = np.zeros(7)
self.do_simulation(u, self.frame_skip)
def get_env_state(self):
joint_state = self.sim.get_state()
mocap_state = self.data.mocap_pos, self.data.mocap_quat
state = joint_state, mocap_state
return copy.deepcopy(state)
def set_env_state(self, state):
joint_state, mocap_state = state
self.sim.set_state(joint_state)
mocap_pos, mocap_quat = mocap_state
self.data.set_mocap_pos('mocap', mocap_pos)
self.data.set_mocap_quat('mocap', mocap_quat)
self.sim.forward()
def take_images(self, filename, obj, writer, img_idx=0, debug=False):
model = load_model_from_path(filename)
sim = MjSim(model)
modder= TextureModder(sim)
# viewer=MjViewer(sim) # this f*****g line has caused me so much pain.
embedding = np.append(obj.type, obj.geom.reshape(-1))
if obj.type == 4 or obj.type == 5:
# MULTI CABINET: get double the params.
axis1, door1 = transform_param(obj.params[0][0],obj.params[0][1], obj)
axis2, door2 = transform_param(obj.params[1][0],obj.params[1][1], obj)
axes = np.append(axis1,axis2)
doors = np.append(door1,door2)
params = np.append(axes,doors)
set_two_door_control(sim, 'cabinet2' if obj.type == 4 else 'refrigerator')
n_qpos_variables = 2
else:
n_qpos_variables = 1
if obj.type == 1:
sim.data.ctrl[0] = 0.05
elif obj.geom[3] == 1:
sim.data.ctrl[0] = -0.2
else:
sim.data.ctrl[0] = 0.2
params = get_cam_relative_params2(obj) # if 1DoF, params is length 10. If 2DoF, params is length 20.
embedding_and_params = np.concatenate((embedding, params, obj.pose, obj.rotation))
# print('nqpos', n_qpos_variables)
# print(self.img_idx, obj.pose)
# print(embedding_and_params.shape)
t = 0
while t<4000:
sim.forward()
sim.step()
if t%250==0:
#########################
IMG_WIDTH = calibrations.sim_width
IMG_HEIGHT = calibrations.sim_height
#########################
img, depth = sim.render(IMG_WIDTH, IMG_HEIGHT, camera_name='external_camera_0', depth=True)
depth = vertical_flip(depth)
real_depth = buffer_to_real(depth, 12.0, 0.1)
norm_depth = real_depth / 12.0
if self.masked:
# remove background
mask = norm_depth > 0.99
norm_depth = (1-mask)*norm_depth
if self.debugging:
# save image to disk for visualization
# img = cv2.resize(img, (IMG_WIDTH,IMG_HEIGHT))
img = vertical_flip(img)
img = white_bg(img)
imgfname = os.path.join(self.savedir, 'img'+str(self.img_idx).zfill(6)+'.png')
depth_imgfname = os.path.join(self.savedir, 'depth_img'+str(self.img_idx).zfill(6)+'.png')
integer_depth = norm_depth * 255
cv2.imwrite(imgfname, img)
cv2.imwrite(depth_imgfname, integer_depth)
# noise_img = sp_noise(img,0.1)
# cv2.imwrite(imgfname, noise_img)
# noise_integer_depth = sp_noise(integer_depth,0.1)
# cv2.imwrite(depth_imgfname, noise_integer_depth)
# if IMG_WIDTH != 192 or IMG_HEIGHT != 108:
# depth = cv2.resize(norm_depth, (192,108))
depthfname = os.path.join(self.savedir,'depth'+str(self.img_idx).zfill(6) + '.pt')
torch.save(torch.tensor(norm_depth.copy()), depthfname)
row = np.append(embedding_and_params, sim.data.qpos[:n_qpos_variables])
# print(row.shape)
writer.writerow(row)
self.img_idx += 1
t += 1
# sim = MjSim(model)
env = gym.make('FetchPickAndPlace-v1')
# exit()
# env.sim = sim
sim = env.sim
viewer = MjRenderContextOffscreen(sim)
viewer.cam.fixedcamid = 3
viewer.cam.type = const.CAMERA_FIXED
env.env._viewers['rgb_array'] = viewer
im = env.render(mode="rgb_array")
plt.imshow(im)
plt.show()
modder = TextureModder(sim)
# modder = CameraModder(sim)
modder.whiten_materials()
# modder = MaterialModder(sim)
t = 1
# viewer.cam.fixedcamid = 3
# viewer.cam.type = const.CAMERA_FIXED
while True:
randomize_textures(sim)
# randomize_lights(sim)
if t % 10 == 0:
randomize_camera(viewer)
def test_textures():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.premod.png')
random_state = np.random.RandomState(0)
modder = TextureModder(sim, random_state=random_state)
modder.whiten_materials()
modder.whiten_materials(['g1', 'g2'])
modder.set_rgb('g1', (255, 0, 0))
modder.set_rgb('g2', (0, 255, 0))
modder.set_rgb('g3', (0, 0, 255))
modder.set_rgb('g4', (255, 0, 255))
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rgb.png')
modder.set_checker('g1', (255, 0, 0), (0, 255, 0))
modder.set_gradient('g2', (0, 255, 0), (0, 0, 255), vertical=True)
modder.set_gradient('g3', (255, 255, 0), (0, 0, 255), vertical=False)
modder.set_noise('g4', (0, 0, 255), (255, 0, 0), 0.1)
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.variety.png')
modder.rand_checker('g1')
modder.rand_gradient('g2')
modder.rand_noise('g3')
modder.rand_rgb('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_specific.png')
modder.rand_all('g1')
modder.rand_all('g2')
modder.rand_all('g3')
modder.rand_all('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_all.png')
modder.rand_checker('g1')
modder.rand_checker('g2')
modder.rand_checker('g3')
modder.rand_checker('g4')
mat_modder = MaterialModder(sim, random_state=random_state)
mat_modder.rand_texrepeat('g1')
mat_modder.rand_texrepeat('g2')
mat_modder.rand_texrepeat('g3')
mat_modder.rand_texrepeat('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_texrepeat.png')
def test_textures():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.premod.png')
random_state = np.random.RandomState(0)
modder = TextureModder(sim, random_state=random_state)
modder.whiten_materials()
modder.whiten_materials(['g1', 'g2'])
modder.set_rgb('g1', (255, 0, 0))
modder.set_rgb('g2', (0, 255, 0))
modder.set_rgb('g3', (0, 0, 255))
modder.set_rgb('g4', (255, 0, 255))
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rgb.png')
modder.set_checker('g1', (255, 0, 0), (0, 255, 0))
modder.set_gradient('g2', (0, 255, 0), (0, 0, 255), vertical=True)
modder.set_gradient('g3', (255, 255, 0), (0, 0, 255), vertical=False)
modder.set_noise('g4', (0, 0, 255), (255, 0, 0), 0.1)
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.variety.png')
modder.rand_checker('g1')
modder.rand_gradient('g2')
modder.rand_noise('g3')
modder.rand_rgb('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_specific.png')
modder.rand_all('g1')
modder.rand_all('g2')
modder.rand_all('g3')
modder.rand_all('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_all.png')
modder.rand_checker('g1')
modder.rand_checker('g2')
modder.rand_checker('g3')
modder.rand_checker('g4')
mat_modder = MaterialModder(sim, random_state=random_state)
mat_modder.rand_texrepeat('g1')
mat_modder.rand_texrepeat('g2')
mat_modder.rand_texrepeat('g3')
mat_modder.rand_texrepeat('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_textures.rand_texrepeat.png')
class BlockWordEnv:
def __init__(self,
env_file,
random_color=None,
random_num=5,
debug=False,
random_seed=0):
self.env_file = env_file
self.random_color = random_color
self.random_num = random_num
self.debug = debug
self.model = load_model_from_path(env_file)
self.sim = MjSim(self.model, nsubsteps=1)
self.sim.model.vis.map.znear = 0.02
self.sim.model.vis.map.zfar = 50.0
self.cube_size = self.sim.model.geom_size[
self.model._geom_name2id['cube_0']]
self.cuboid_size = self.sim.model.geom_size[
self.model._geom_name2id['cuboid_0']]
self.cube_num = len(
[i for i in self.sim.model.geom_names if 'cube_' in i])
self.cuboid_num = len(
[i for i in self.sim.model.geom_names if 'cuboid_' in i])
if self.debug:
print('total cube num:', self.cube_num)
print('total cuboid num:', self.cuboid_num)
self.max_num_per_type = max(self.cube_num, self.cuboid_num)
self.center_bounds = [
-0.25, 0.25
] # [0, self.cuboid_size[0] * self.max_num_per_type]
self.pos_candidates = np.arange(
self.center_bounds[0], self.center_bounds[1] + self.cube_size[0],
self.cube_size[0])
self.modder = TextureModder(self.sim, random_state=random_seed)
self.cur_id = {'cube': 0, 'cuboid': 0}
self.viewer = None
self.active_blocks = []
if random_color:
self.reset_viewer()
def reset(self):
self.sim.reset()
self.sim.step()
self.active_blocks = []
self.cur_id = {'cube': 0, 'cuboid': 0}
if self.random_color:
self.randomize_color()
if self.viewer is not None:
self.reset_viewer()
def randomize_color(self):
self.modder.whiten_materials()
for name in self.sim.model.geom_names:
if 'table' in name:
continue
self.modder.rand_all(name)
def simulate_one_epoch(self):
self.gen_constrained_ran_bk_configs()
imgs = []
for i in range(self.random_num):
img = self.get_img()
imgs.append(img.copy())
self.randomize_color()
stable = self.check_stability(render=False)
return imgs, stable
def step(self):
self.sim.step()
def forward(self):
self.sim.forward()
def get_active_bk_states(self):
bk_poses = []
for bk_name in self.active_blocks:
pose = self.sim.data.get_joint_qpos(bk_name)
bk_poses.append(pose)
return np.array(bk_poses)
def check_stability(self, render=False):
self.sim.step()
prev_poses = self.get_active_bk_states()
for i in range(400):
self.sim.step()
if render:
self.render()
post_poses = self.get_active_bk_states()
diff = np.abs(post_poses - prev_poses)
diff_norm = np.linalg.norm(diff[:, :3], axis=1)
if self.debug:
print('prev:')
print(prev_poses[:, :3])
print('post')
print(post_poses[:, :3])
print('diff norm:', diff_norm)
stable = np.all(diff_norm < 0.01)
if self.debug:
print('Current configuration stable:', stable)
return stable
def render(self):
if self.viewer is None:
self.reset_viewer()
self.viewer.render()
def reset_viewer(self):
self.viewer = MjViewer(self.sim)
self.viewer.cam.lookat[:3] = np.array([0, 0, 0.1])
self.viewer.cam.distance = 2
self.viewer.cam.elevation = -20
def move_given_block(self, name, target_pos):
prev_pose = self.sim.data.get_joint_qpos(name)
post_pose = prev_pose.copy()
post_pose[:3] = target_pos
self.sim.data.set_joint_qpos(name, post_pose)
self.active_blocks.append(name)
if self.debug:
print('{0:s} pose after moving:'.format(name), post_pose)
def move_given_blocks(self, block_dict):
for bk_name, pos in block_dict.items():
self.move_given_block(bk_name, pos)
def move_block_for_demo(self, name, target_pos):
prev_pose = self.sim.data.get_joint_qpos(name)
if self.debug:
print('{0:s} pose before moving:'.format(name), prev_pose)
post_pose = prev_pose.copy()
planned_path = []
up_steps = 20
h_steps = 30
down_steps = 20
planned_path.append(prev_pose.copy())
for i in range(up_steps):
tmp_pose = planned_path[-1].copy()
tmp_pose[2] += (0.45 - prev_pose[2]) / float(up_steps)
planned_path.append(tmp_pose.copy())
for i in range(h_steps + 1):
tmp_pose = planned_path[-1].copy()
tmp_pose[0] = (target_pos[0] -
prev_pose[0]) * i / h_steps + prev_pose[0]
tmp_pose[1] = (target_pos[1] -
prev_pose[1]) * i / h_steps + prev_pose[1]
planned_path.append(tmp_pose.copy())
for i in range(down_steps):
tmp_pose = planned_path[-1].copy()
tmp_pose[2] -= (0.45 - target_pos[2]) / float(down_steps)
planned_path.append(tmp_pose.copy())
post_pose[:3] = target_pos
planned_path.append(post_pose.copy())
for pos in planned_path:
self.sim.data.set_joint_qpos(name, pos)
time.sleep(0.02)
self.sim.forward()
self.render()
self.active_blocks.append(name)
if self.debug:
print('{0:s} pose after moving:'.format(name), post_pose)
def move_blocks_for_demo(self, block_dict):
name = list(block_dict.keys())[0]
target_pos = block_dict[name]
initial_poses = {}
for key in block_dict.keys():
initial_poses[key] = self.sim.data.get_joint_qpos(key)
prev_pose = self.sim.data.get_joint_qpos(name)
if self.debug:
print('{0:s} pose before moving:'.format(name), prev_pose)
post_pose = prev_pose.copy()
planned_delta_path = []
planned_path = []
up_steps = 20
h_steps = 30
down_steps = 20
planned_path.append(prev_pose.copy())
planned_delta_path.append(np.zeros_like(prev_pose))
for i in range(up_steps):
tmp_pose = planned_path[-1].copy()
tmp_pose[2] += (0.45 - prev_pose[2]) / float(up_steps)
planned_delta_path.append(tmp_pose - planned_path[-1])
planned_path.append(tmp_pose.copy())
for i in range(h_steps + 1):
tmp_pose = planned_path[-1].copy()
tmp_pose[0] = (target_pos[0] -
prev_pose[0]) * i / h_steps + prev_pose[0]
tmp_pose[1] = (target_pos[1] -
prev_pose[1]) * i / h_steps + prev_pose[1]
planned_delta_path.append(tmp_pose - planned_path[-1])
planned_path.append(tmp_pose.copy())
for i in range(down_steps):
tmp_pose = planned_path[-1].copy()
tmp_pose[2] -= (0.45 - target_pos[2]) / float(down_steps)
planned_delta_path.append(tmp_pose - planned_path[-1])
planned_path.append(tmp_pose.copy())
post_pose[:3] = target_pos
planned_delta_path.append(post_pose - planned_path[-1])
planned_path.append(post_pose.copy())
for delta_pos in planned_delta_path:
for bk_name, target_bk_pos in block_dict.items():
initial_poses[bk_name] = initial_poses[bk_name] + delta_pos
self.sim.data.set_joint_qpos(bk_name, initial_poses[bk_name])
time.sleep(0.02)
self.sim.forward()
self.render()
self.active_blocks.append(name)
if self.debug:
print('{0:s} pose after moving:'.format(name),
self.sim.data.get_joint_qpos(name))
def move_block(self, target_pos, bk_type='cube'):
# center bounds: [0, 0.1 * 30]
assert bk_type == 'cube' or bk_type == 'cuboid'
bk_name = '{0:s}_{1:d}'.format(bk_type, self.cur_id[bk_type])
prev_pose = self.sim.data.get_joint_qpos(bk_name)
post_pose = prev_pose.copy()
post_pose[:3] = target_pos
self.sim.data.set_joint_qpos(bk_name, post_pose)
self.active_blocks.append(bk_name)
if self.debug:
print(
'{0:s}_{1:d} pose after moving:'.format(
bk_type, self.cur_id[bk_type]), post_pose)
self.cur_id[bk_type] += 1
def get_img(self):
# return self.get_img_demo()
img = self.sim.render(camera_name='camera',
width=600,
height=600,
depth=False)
img = np.flipud(img)
resized_img = cv2.resize(img[0:500, 50:550], (224, 224),
cv2.INTER_AREA)
return resized_img
def get_img_demo(self):
img = self.sim.render(camera_name='camera',
width=1200,
height=700,
depth=False)
img = np.flipud(img)
img = img[:, :, ::-1]
return img
def build_tower(self):
layer_num = 1
for i in range(20):
z = (2 * layer_num - 1) * self.cube_size[2]
y = 0
x = 0
target_pos = np.array([x, y, z])
self.move_block(target_pos, bk_type='cube')
layer_num += 1
def gen_ran_bk_configs(self, render=False):
while True:
cuboid_num = np.random.choice(5, 1)[0]
cube_num = np.random.choice(15, 1)[0]
if cuboid_num > 0 or cube_num > 0:
break
total_num = cuboid_num + cube_num
blocks = [0] * cube_num + [1] * cuboid_num
permuted_blocks = np.random.permutation(blocks)
cur_x = self.center_bounds[0]
layer_num = 1
for i in range(total_num):
bk = permuted_blocks[i]
bk_type = 'cube' if bk == 0 else 'cuboid'
bk_size = self.cube_size if bk == 0 else self.cuboid_size
z = (2 * layer_num - 1) * self.cube_size[2]
y = 0
if self.center_bounds[1] - cur_x < bk_size[0]:
layer_num += 1
cur_x = self.center_bounds[0]
continue
else:
bk_lower_limit = cur_x + bk_size[0]
pos_candidates = self.pos_candidates[
self.pos_candidates >= bk_lower_limit]
x = np.random.choice(pos_candidates, 1)[0]
cur_x = x + bk_size[0]
target_pos = np.array([x, y, z])
self.move_block(target_pos, bk_type=bk_type)
self.sim.forward()
if render:
self.render()
def gen_constrained_ran_bk_configs(self, render=False):
# prob = np.exp(-0.1 * np.arange(30))
while True:
cuboid_num = np.random.choice(5, 1)[0]
cube_num = np.random.choice(15, 1)[0]
if cuboid_num > 0 or cube_num > 0:
break
if self.debug:
print('Selected cube num:', cube_num)
print('Selected cuboid num:', cuboid_num)
total_num = cuboid_num + cube_num
blocks = [0] * cube_num + [1] * cuboid_num
permuted_blocks = np.random.permutation(blocks)
cur_x = self.center_bounds[0]
layer_num = 1
layer_pos_candidates = self.pos_candidates.copy()
filled_segs = []
for i in range(total_num):
bk = permuted_blocks[i]
bk_type = 'cube' if bk == 0 else 'cuboid'
bk_size = self.cube_size if bk == 0 else self.cuboid_size
z = (2 * layer_num - 1) * self.cube_size[2]
y = 0
bk_lower_limit = cur_x + bk_size[0]
pos_candidates = layer_pos_candidates[
layer_pos_candidates >= bk_lower_limit]
if pos_candidates.size < 1:
layer_num += 1
cur_x = self.center_bounds[0]
layer_pos_candidates = self.pos_candidates.copy()
good_ids = np.zeros_like(layer_pos_candidates, dtype=bool)
for seg in filled_segs:
good_ids = np.logical_or(
good_ids,
np.logical_and(layer_pos_candidates >= seg[0],
layer_pos_candidates <= seg[1]))
layer_pos_candidates = layer_pos_candidates[good_ids]
if self.debug:
print('Layer [{0:d}] pos candidates num: {1:d}'.format(
layer_num, layer_pos_candidates.size))
if layer_pos_candidates.size < 1:
break
filled_segs = []
continue
else:
x = np.random.choice(pos_candidates, 1)[0]
cur_x = x + bk_size[0]
target_pos = np.array([x, y, z])
self.move_block(target_pos, bk_type=bk_type)
filled_segs.append([x - bk_size[0], cur_x])
self.sim.forward()
if render:
self.render()
def render_callback(sim, viewer):
global modder
if modder is None:
modder = TextureModder(sim)
for name in sim.model.geom_names:
modder.rand_all(name)
import numpy as np
from mujoco_py import load_model_from_path, MjSim, MjViewer
from mujoco_py.modder import TextureModder
env_file = 'block_world.xml'
model = load_model_from_path(env_file)
sim = MjSim(model, nsubsteps=1)
sim.model.vis.map.znear = 0.02
sim.model.vis.map.zfar = 50.0
modder = TextureModder(sim)
viewer = MjViewer(sim)
viewer.cam.lookat[:3] = np.array([0, 0, 0.1])
viewer.cam.distance = 2
viewer.cam.elevation = -20
# sim.model.body_pos[1] = np.array([0.6, 0.6, 4])
sim.reset()
cube_pose = sim.data.get_joint_qpos('cube_0')
cube_pose[:3] = np.array([0, 0, 0.02])
sim.data.set_joint_qpos('cube_0', cube_pose)
cube_pose = sim.data.get_joint_qpos('cuboid_0')
cube_pose[:3] = np.array([0., 0, 0.06])
sim.data.set_joint_qpos('cuboid_0', cube_pose)
cube_pose = sim.data.get_joint_qpos('cube_1')
cube_pose[:3] = np.array([0.12, 0, 0.02])
sim.data.set_joint_qpos('cube_1', cube_pose)
cube_pose = sim.data.get_joint_qpos('cube_2')
cube_pose[:3] = np.array([0.08, 0, 0.10])
sim.data.set_joint_qpos('cube_2', cube_pose)
#!/usr/bin/env python3
"""
Displays robot fetch at a disco party.
"""
from mujoco_py import load_model_from_path, MjSim, MjViewer
from mujoco_py.modder import TextureModder
import os
model = load_model_from_path("xmls/fetch/main.xml")
sim = MjSim(model)
viewer = MjViewer(sim)
modder = TextureModder(sim)
t = 0
while True:
for name in sim.model.geom_names:
modder.rand_all(name)
viewer.render()
t += 1
if t > 100 and os.getenv('TESTING') is not None:
break
"""
Simulation of Nasa Hand
"""
from mujoco_py import load_model_from_path, MjSim, MjViewer
from mujoco_py.modder import TextureModder
import hand_sim_mj.utils.util_parser_disp as util
import os
import time
# constants
# load the dynamic model
model_xml_path = os.path.join(os.environ[util.SIM_ENV_VAR],
'xml/nasa_hand/main.xml')
model = load_model_from_path(model_xml_path)
sim = MjSim(model)
viewer = MjViewer(sim)
modder = TextureModder(sim)
# set initial joint valuse
# sim.data.set_joint_qpos('palm_slide_X', 0.001)
# sim.data.set_joint_qpos('palm_slide_Y', 0.001)
# sim.data.set_joint_qpos('palm_slide_Z', 0.001)
# A = PyKDL.Rotation.RotX(0/180*util.PI)
# q = A.GetQuaternion()
# sim.data.set_joint_qpos('palm_ball_joint', q)
sim.data.set_joint_qpos('finger1_roll_joint', 1.0)
sim.data.set_joint_qpos('finger1_prox_joint', 0.0)
sim.data.set_joint_qpos('finger1_dist_joint', 0.0)
sim.data.set_joint_qpos('finger2_roll_joint', -1.0)
sim.data.set_joint_qpos('finger2_prox_joint', 0.0)
sim.data.set_joint_qpos('finger2_dist_joint', 0.0)
sim.data.set_joint_qpos('thumb_prox_joint', 0.0)
sim.data.set_joint_qpos('thumb_dist_joint', 0.0)
"""
Displays robot fetch at a disco party.
"""
from mujoco_py import load_model_from_path, MjSim, MjViewer
from mujoco_py.modder import TextureModder, MaterialModder
import os
from gym import utils
from gym.envs.robotics import gen3_env
MODEL_XML_PATH = "/home/rjangir/workSpace/gen3-mujoco/gym/gym/envs/robotics/assets/gen3/sideways_fold.xml"
model = load_model_from_path(MODEL_XML_PATH)
#model = load_model_from_path("xmls/fetch/main.xml")
sim = MjSim(model)
viewer = MjViewer(sim)
modder = TextureModder(sim)
#matModder = MaterialModder(sim) #specularity, shininess, reflectance
t = 0
while True:
#print("body names", sim.model.name_skinadr, sim.model.skin_matid[0])
skin_mat = sim.model.skin_matid[0]
#print("skin mat texture", sim.model.mat_texid[skin_mat], sim.model.tex_type[sim.model.mat_texid[skin_mat]])
for name in sim.model.geom_names:
modder.whiten_materials()
modder.set_checker(name, (255, 0, 0), (0, 0, 0))
modder.rand_all(name)
modder.set_checker('skin', (255, 0, 0), (0, 0, 0))
modder.rand_all('skin')
viewer.render()
class BlocksEnv(robot_env.RobotEnv):
"""Superclass for all Fetch environments.
"""
def __init__(
self, model_path, n_substeps, gripper_extra_height, block_gripper,
has_object, target_in_the_air, target_offset, obj_range, target_range,
distance_threshold, initial_qpos, reward_type, num_blocks
):
"""Initializes a new Fetch environment.
Args:
model_path (string): path to the environments XML file
n_substeps (int): number of substeps the simulation runs on every call to step
gripper_extra_height (float): additional height above the table when positioning the gripper
block_gripper (boolean): whether or not the gripper is blocked (i.e. not movable) or not
has_object (boolean): whether or not the environment has an object
target_in_the_air (boolean): whether or not the target should be in the air above the table or on the table surface
target_offset (float or array with 3 elements): offset of the target
obj_range (float): range of a uniform distribution for sampling initial object positions
target_range (float): range of a uniform distribution for sampling a target
distance_threshold (float): the threshold after which a goal is considered achieved
initial_qpos (dict): a dictionary of joint names and values that define the initial configuration
reward_type ('sparse' or 'dense'): the reward type, i.e. sparse or dense
"""
self.gripper_extra_height = gripper_extra_height
self.block_gripper = block_gripper
self.target_in_the_air = target_in_the_air
self.target_offset = target_offset
self.obj_range = obj_range
self.target_range = target_range
self.distance_threshold = distance_threshold
# This should always be sparse
self.reward_type = reward_type
# The number of objects present in the scene
# Now this number is fixed
self.num_objs = num_blocks + 2
self.obj_colors = self._sample_colors()
# This is the goal vector that we define for this task
self.achieved_goal = -1 * np.ones([self.num_objs, self.num_objs])
self.has_succeeded = False
self.id2obj = None
# For curriculum learning
self.difficulty = 0
super(BlocksEnv, self).__init__(
model_path=model_path, n_substeps=n_substeps, n_actions=4,
initial_qpos=initial_qpos)
def _sample_from_table(self):
return np.asarray([TABLE_X + self.np_random.uniform(-TABLE_W, TABLE_W),
TABLE_Y + self.np_random.uniform(-TABLE_H, TABLE_H)])
# For curriculum learning
def increase_difficulty(self):
raise NotImplementedError()
def get_difficulty(self):
return self.difficulty
# Configure the environment for testing
def set_test(self):
raise NotImplementedError()
def _sample_colors(self):
raise NotImplementedError()
def _geom2objid(self, i):
name = self.sim.model.geom_id2name(i)
# First object is reserved for the gripper
# and the second object is reserved for the table
if name != None:
if "finger" in name:
return 0
elif name == "table":
return 1
elif "object" in name:
return int(name[6:]) + 2
return None
def _check_goal(self):
for i in range(self.num_objs):
for j in range(self.num_objs):
if (self.achieved_goal[i][j] != self.achieved_goal[j][i]):
return False
return True
# GoalEnv methods
# ----------------------------
# Goal is defined as a matrix of touching conditions
# -1 for never touched, 1 for is touching, 0 for touched but not touching
# In the desired goal we simply input 1 for the blocks we want touching
# and -1 for blocks we never want touch, and 0 for the blocks we don't care
# In this setting, if the achieved goal aligns with the desired goal
# their dot product should equal to the number of non-zero elements
def compute_reward(self, achieved_goal, goal, info):
# Compute distance between goal and the achieved goal.
# goal_len = min(achieved_goal.shape[0], goal.shape[0])
# achieved_goal = achieved_goal[:goal_len]
# goal = goal[:goal_len]
d = np.sum(achieved_goal * goal, axis=-1)
c = np.count_nonzero(goal, axis=-1)
return -(d != c).astype(np.float32)
# RobotEnv methods
# ----------------------------
def _step_callback(self):
if self.block_gripper:
self.sim.data.set_joint_qpos('robot0:l_gripper_finger_joint', 0.)
self.sim.data.set_joint_qpos('robot0:r_gripper_finger_joint', 0.)
self.sim.forward()
# Update achieved goal
for i in range(self.num_objs):
for j in range(self.num_objs):
if (self.achieved_goal[i][j] == 1):
self.achieved_goal[i][j] = 0
d = self.sim.data
for i in range(d.ncon):
con = d.contact[i]
i1 = con.geom1
i2 = con.geom2
obj1 = self.id2obj[i1]
obj2 = self.id2obj[i2]
if (obj1 != None and obj2 != None):
self.achieved_goal[obj1][obj2] = 1
self.achieved_goal[obj2][obj1] = 1
# No changes here
def _set_action(self, action):
assert action.shape == (4,)
action = action.copy() # ensure that we don't change the action outside of this scope
pos_ctrl, gripper_ctrl = action[:3], action[3]
pos_ctrl *= 0.05 # limit maximum change in position
rot_ctrl = [1., 0., 1., 0.] # fixed rotation of the end effector, expressed as a quaternion
gripper_ctrl = np.array([gripper_ctrl, gripper_ctrl])
assert gripper_ctrl.shape == (2,)
if self.block_gripper:
gripper_ctrl = np.zeros_like(gripper_ctrl)
action = np.concatenate([pos_ctrl, rot_ctrl, gripper_ctrl])
# Apply action to simulation.
utils.ctrl_set_action(self.sim, action)
utils.mocap_set_action(self.sim, action)
def _get_obs(self):
# positions
grip_pos = self.sim.data.get_site_xpos('robot0:grip')
dt = self.sim.nsubsteps * self.sim.model.opt.timestep
grip_velp = self.sim.data.get_site_xvelp('robot0:grip') * dt
robot_qpos, robot_qvel = utils.robot_get_obs(self.sim)
object_pos, object_rot, object_velp, object_velr, object_rel_pos = [], [], [], [], []
gripper_state = robot_qpos[-2:]
gripper_vel = robot_qvel[-2:] * dt # change to a scalar if the gripper is made symmetric
obs = []
for i in range(self.num_objs-2):
obj_name = 'object{}'.format(i)
temp_pos = self.sim.data.get_site_xpos(obj_name)
# rotations
temp_rot = rotations.mat2euler(self.sim.data.get_site_xmat(obj_name))
# velocities
temp_velp = self.sim.data.get_site_xvelp(obj_name) * dt
temp_velr = self.sim.data.get_site_xvelr(obj_name) * dt
# gripper state
temp_rel_pos = temp_pos - grip_pos
temp_velp -= grip_velp
block_obs = np.concatenate([temp_pos.ravel(),
temp_rel_pos.ravel(), temp_rot.ravel(),
temp_velp.ravel(), temp_velr.ravel()])
obs = np.concatenate([obs, block_obs])
assert(self._check_goal())
achieved_goal = self.achieved_goal.copy().ravel()
obs = np.concatenate([grip_pos, gripper_state, grip_velp, gripper_vel, obs])
return {
'observation': obs.copy(),
'achieved_goal': achieved_goal.copy(),
'desired_goal': self.goal.copy(),
}
def _viewer_setup(self):
body_id = self.sim.model.body_name2id('robot0:gripper_link')
lookat = self.sim.data.body_xpos[body_id]
for idx, value in enumerate(lookat):
self.viewer.cam.lookat[idx] = value
self.viewer.cam.distance = 2.5
self.viewer.cam.azimuth = 132.
self.viewer.cam.elevation = -14.
def _render_callback(self):
# Visualize target.
# sites_offset = (self.sim.data.site_xpos - self.sim.model.site_pos).copy()
# site_id = self.sim.model.site_name2id('target0')
# self.sim.model.site_pos[site_id] = self.goal - sites_offset[0]
# self.sim.forward()
pass
def _reset_sim(self, test=False):
self.sim.set_state(self.initial_state)
# Randomize colors for each object
self.obj_colors = self._sample_colors()
# Randomize start position of each object
self._randomize_objects(test)
self.sim.forward()
self.has_succeeded = False
return True
def _randomize_objects(self, test=False):
raise NotImplementedError()
def _sample_goal(self):
C = self.obj_colors
goal = []
for i in range(self.num_objs):
for j in range(self.num_objs):
if ((C[i] == RED and C[j] == BLUE) or
(C[i] == BLUE and C[j] == RED)):
goal.append(-1)
elif ((C[i] == GREEN and C[j] == BLUE) or
(C[i] == BLUE and C[j] == GREEN)):
goal.append(1)
else:
goal.append(0)
return np.asarray(goal)
def _is_success(self, achieved_goal, desired_goal):
# Override whatever input it gives
achieved_goal, desired_goal = self.achieved_goal.ravel(), self.goal
r = self.compute_reward(achieved_goal, desired_goal, None)
if r == 0:
self.has_succeeded = True
return self.has_succeeded
def _env_setup(self, initial_qpos):
self.id2obj = [self._geom2objid(i) for i in range(self.sim.model.ngeom)]
for name, value in initial_qpos.items():
self.sim.data.set_joint_qpos(name, value)
utils.reset_mocap_welds(self.sim)
self.sim.forward()
# Move end effector into position.
gripper_target = np.array([-0.498, 0.005, -0.431 + self.gripper_extra_height]) + self.sim.data.get_site_xpos('robot0:grip')
gripper_rotation = np.array([1., 0., 1., 0.])
self.sim.data.set_mocap_pos('robot0:mocap', gripper_target)
self.sim.data.set_mocap_quat('robot0:mocap', gripper_rotation)
for _ in range(10):
self.sim.step()
# Extract information for sampling goals.
if not hasattr(self, 'initial_gripper_xpos'):
self.initial_gripper_xpos = self.sim.data.get_site_xpos('robot0:grip').copy()
self.height_offset = self.sim.data.get_site_xpos('object0')[2]
# Change the colors to match the success conditions
self.color_modder = TextureModder(self.sim)
for i in range(self.sim.model.ngeom):
obj_id = self.id2obj[i]
if obj_id != None:
name = self.sim.model.geom_id2name(i)
if 'table' in name:
color = np.asarray(COLORS_RGB[self.obj_colors[obj_id]])
self.color_modder.set_rgb(name, color * 2)
else:
self.color_modder.set_rgb(name,
COLORS_RGB[self.obj_colors[obj_id]])
def render_callback(sim, viewer):
global modder
if modder is None:
modder = TextureModder(sim)
for name in sim.model.geom_names:
modder.rand_all(name)