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(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 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')
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
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 render_callback(sim, viewer): global modder if modder is None: modder = TextureModder(sim) for name in sim.model.geom_names: modder.rand_all(name)
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)
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) sim.forward() cam_img = sim.render(640, 360)[::-1, :, :] import ipdb; ipdb.set_trace()
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() t += 1 if t > 100 and os.getenv('TESTING') is not None: break
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()
#!/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
class RobotEnv(gym.GoalEnv): def __init__(self, model_path, initial_qpos, n_actions, n_substeps, mode, visual_randomize, visual_data_recording): 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)) self.mode = mode 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._viewers = {} self.modder = TextureModder(self.sim) self.visual_randomize = visual_randomize 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.visual_data_recording = visual_data_recording 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() @property def dt(self): return self.sim.model.opt.timestep * self.sim.nsubsteps # Env methods # ---------------------------- def seed(self, seed=None): self.np_random, seed = seeding.np_random(seed) return [seed] def step(self, action): if self.visual_randomize: for name in self.sim.model.geom_names: self.modder.whiten_materials() self.modder.set_checker(name, (255, 0, 0), (0, 0, 0)) self.modder.rand_all(name) self.mat_modder.rand_all(name) self.modder.set_checker('skin', (255, 0, 0), (0, 0, 0)) self.modder.rand_all('skin') self.light_modder.rand_all('light0') self.light_modder.rand_all('light1') self.light_modder.rand_all('light2') self.light_modder.rand_all('light3') self.camera_modder.rand_fovy('camera1') self.camera_modder.rand_pos('camera1') self.camera_modder.rand_quat('camera1', factor=[0.1, 0.1, 0.1]) action = np.clip(action, self.action_space.low, self.action_space.high) self._set_action(action) self.sim.step() self._step_callback() obs = self._get_obs() done = False info = { 'is_success': self._is_success(obs['achieved_goal'], self.goal), } reward = self.compute_reward(obs['achieved_goal'], self.goal, info) #print("REWARD", reward ," and IS SUCCESS ", info['is_success']) return obs, reward, done, info def reset(self): # Attempt to reset the simulator. Since we randomize initial conditions, it # is possible to get into a state with numerical issues (e.g. due to penetration or # Gimbel lock) or we may not achieve an initial condition (e.g. an object is within the hand). # In this case, we just keep randomizing until we eventually achieve a valid initial # configuration. super(RobotEnv, self).reset() did_reset_sim = False while not did_reset_sim: did_reset_sim = self._reset_sim() # if self.visual_randomize: # for name in self.sim.model.geom_names: # self.modder.whiten_materials() # self.modder.set_checker(name, (255, 0, 0), (0, 0, 0)) # self.modder.rand_all(name) # self.mat_modder.rand_all(name) # self.modder.set_checker('skin', (255, 0, 0), (0, 0, 0)) # self.modder.rand_all('skin') self.goal = self._sample_goal().copy() obs = self._get_obs() return obs def close(self): if self.viewer is not None: # self.viewer.finish() self.viewer = None self._viewers = {} def render(self, mode='human', width=DEFAULT_SIZE, height=DEFAULT_SIZE): self._render_callback() if mode == 'rgb_array': self._get_viewer(mode).render(width, height) # window size used for old mujoco-py: data = self._get_viewer(mode).read_pixels(width, height, depth=False) # original image is upside-down, so flip it return data[::-1, :, :] #return None elif mode == 'human': self._get_viewer(mode).render() def _get_viewer(self, mode): self.viewer = self._viewers.get(mode) if self.viewer is None: if mode == 'human': self.viewer = mujoco_py.MjViewer(self.sim) elif mode == 'rgb_array': device_id = self.sim.model.camera_name2id('camera1') self.viewer = mujoco_py.MjRenderContextOffscreen(self.sim, device_id=device_id) self._viewer_setup() self._viewers[mode] = self.viewer return self.viewer # Extension methods # ---------------------------- def _reset_sim(self): """Resets a simulation and indicates whether or not it was successful. If a reset was unsuccessful (e.g. if a randomized state caused an error in the simulation), this method should indicate such a failure by returning False. In such a case, this method will be called again to attempt a the reset again. """ self.sim.set_state(self.initial_state) self.sim.forward() return True def _get_obs(self): """Returns the observation. """ raise NotImplementedError() def _set_action(self, action): """Applies the given action to the simulation. """ raise NotImplementedError() def _is_success(self, achieved_goal, desired_goal): """Indicates whether or not the achieved goal successfully achieved the desired goal. """ raise NotImplementedError() def _sample_goal(self): """Samples a new goal and returns it. """ raise NotImplementedError() def _env_setup(self, initial_qpos): """Initial configuration of the environment. Can be used to configure initial state and extract information from the simulation. """ pass def _viewer_setup(self): """Initial configuration of the viewer. Can be used to set the camera position, for example. """ pass def _render_callback(self): """A custom callback that is called before rendering. Can be used to implement custom visualizations. """ pass def _step_callback(self): """A custom callback that is called after stepping the simulation. Can be used to enforce additional constraints on the simulation state. """ pass
class MjDataSetGenerator(DataSetGenerator): IMG_HEIGHT = 480 IMG_WIDTH = 640 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 on_screen_render(self, cam_name): self.sim.reset() temp_viewer = mujoco_py.MjViewer(self.sim) t = 0 while True: # Randomised material/texture if any is assigned to geom self._rand_mat() # Set camera position and orientation self._set_cam_pos(cam_name, t) self.sim.step() self._set_cam_orientation(cam_name, t) temp_viewer.render() t += 1 if t > 100 and os.getenv('TESTING') is not None: break def create_data_set(self, ndata, radius_range, deg_range, quat, cameras, target_geom, png_tex_ids=None, start=0): self.sim.reset() self._make_dir() t = start # initialise the camera position array self.cam_pos = self._get_cam_pos(radius_range, deg_range, quat, ndata, start) # self.cam_pos = self._get_debug_cam_pos(radius_range, deg_range, quat, ndata, 4) # generate dataset while True: # Randomised the position of the object # Set camera position self._set_cam_pos(cameras[0], t) # Randomised light source position self._randomise_light_pos() # Randomised material/texture if any is assigned to geom self._rand_mat(target_geom, png_tex_ids) # Simulate and render in offscreen renderer self.sim.step() # Save images for all camera for cam in cameras: self._set_cam_orientation(cam, t) cam_id = self.cam_modder.get_camid(cam) self.viewer.render(self.IMG_WIDTH, self.IMG_HEIGHT, cam_id) rgb = self.viewer.read_pixels(self.IMG_WIDTH, self.IMG_HEIGHT)[0][::-1, :, :] self._save_fig_to_dir(rgb, t, cam_id) # Time advance one t += 1 # Print progress to terminal self.print_progress(ndata, t) if t == ndata or os.getenv('TESTING') is not None: print("Finish creating {} {} images".format( ndata, self.data_set_name)) break self._save_cam_pos(self.cam_pos) def _rand_mat(self, target_geom, png_tex): for name in self.sim.model.geom_names: # Random mat using mujoco py modder or png in xml if self.use_procedural: self.tex_modder.rand_all(name) else: # change redness if name == target_geom["cube"]: self._change_redness(name) # change wood texture if name == target_geom["ground"]: self._change_tex_png(name, png_tex) def _get_tex_id(self, mat_id): tex_id = self.model.mat_texid[mat_id] assert tex_id >= 0, "Material has no assigned texture" return tex_id def _change_redness(self, name): r = self.tex_modder.random_state.randint(120, 256) g = self.tex_modder.random_state.randint(0, 120) b = self.tex_modder.random_state.randint(0, 120) self.tex_modder.set_rgb(name, [r, g, b]) def _change_tex_png(self, name, png_tex): geom_id = self.model.geom_name2id(name) mat_id = self.model.geom_matid[geom_id] self.model.mat_texid[mat_id] = \ self.tex_modder.random_state.randint(png_tex[0], png_tex[1] + 1) self.tex_modder.upload_texture(name) def _set_cam_pos(self, cam_name, t, printPos=None): # If no if self.cam_pos is None: self.cam_pos = self._get_cam_pos([0.25, 0.7], [0, 80], 0.5) # set position of the reference camera cam_id = self.cam_modder.get_camid(cam_name) self.model.cam_pos[cam_id] = self.cam_pos[t, 0:3] if printPos: print("The cam pos is: ", self.cam_pos[t, :]) # Call after sim.step if want to change the camera orientation while keep def _set_cam_orientation(self, cam_name, t, printPos=None): cam_id = self.cam_modder.get_camid(cam_name) if self.cam_pos_file is None: # Add the offset to the rotational matrix self.sim.data.cam_xmat[ cam_id] = self.sim.data.cam_xmat[cam_id] + self.cam_pos[t, 3:] # Save the actual rotational matrix self.cam_pos[t, 3:] = self.sim.data.cam_xmat[cam_id] else: self.sim.data.cam_xmat[cam_id] = self.cam_pos[t, 3:] if printPos: print("The cam orientation is: ", self.cam_pos[t, :]) def _randomise_light_pos(self): # set position # index of light is hard coded for now # TODO: get light index by name for i in range(len(self.model.light_pos)): self.model.light_pos[i, 0] = uniform(-10, 10) self.model.light_pos[i, 1] = uniform(-10, 5)