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 __init__(self, filepath, random_params={}, gpu_render=False, gui=False, display_data=False):
self.model = load_model_from_path(filepath)
self.sim = MjSim(self.model)
self.filepath = filepath
self.gui = gui
self.display_data = display_data
# Take the default random params and update anything we need
self.RANDOM_PARAMS = {}
self.RANDOM_PARAMS.update(random_params)
if gpu_render:
self.viewer = MjViewer(self.sim)
else:
self.viewer = None
# Get start state of params to slightly jitter later
self.START_GEOM_POS = self.model.geom_pos.copy()
self.START_GEOM_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.tex_modder = TextureModder(self.sim)
self.tex_modder.whiten_materials() # ensures materials won't impact colors
self.cam_modder = CameraModder(self.sim)
self.light_modder = LightModder(self.sim)
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 _render_seg(sim: mujoco_py.MjSim, camera: str, render_height: int, render_width: int, world_xml: ET.Element):
lm = LightModder(sim)
# switch all lights off
light_names = [l.attrib['name'] for l in world_xml.findall(".//light")]
for light_name in light_names:
lm.set_active(light_name, 0)
# take screenshot
frame = sim.render(render_width * 2, render_height, camera_name=camera)
# reset lights
for light_name in light_names:
lm.set_active(light_name, 1)
return frame
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 _render_rgb(sim: mujoco_py.MjSim, camera: str, render_height: int,
render_width: int, world_xml: ET.Element):
lm = LightModder(sim)
cam_names = [c.attrib['name'] for c in world_xml.findall(".//camera")]
# lights off for all other cameras except the recording one
for cam in cam_names:
light_name = _get_cam_light_name(cam)
if light_name in sim.model.light_names:
lm.set_active(light_name, 1 if cam == camera else 0)
# take screenshot
frame = sim.render(render_width * 2, render_height, camera_name=camera)
# reset camera lights
for cam in cam_names:
light_name = _get_cam_light_name(cam)
if light_name in sim.model.light_names:
lm.set_active(light_name, 1)
return frame
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 __init__(self, FLAGS):
super().__init__(FLAGS)
if self.FLAGS['baxter']:
self.robot_offset = lambda: self.model.body_pos[self.name2bid(
'base_ground')]
else:
self.robot_offset = lambda: self.model.body_pos[self.name2bid(
'robot_table_floor')]
self.table_center = lambda: TABLES[self.FLAGS['default_table']][
'pos'] + self.robot_offset(
) + [0, 0, TABLES[self.FLAGS['default_table']]['wood'][2]]
self.LIGHT = {
'ambient': self.model.light_ambient,
'active': self.model.light_active,
'castshadow': self.model.light_castshadow,
'diffuse': self.model.light_diffuse,
'dir': self.model.light_dir,
'pos': self.model.light_pos,
'specular': self.model.light_specular
}
self.START_BODY_POS = self.model.body_pos.copy()
self.START_LIGHT = {}
for key in self.LIGHT:
self.START_LIGHT[key] = self.LIGHT[key].copy()
self.tex_modder = TextureModder(self.sim)
self.cam_modder = CameraModder(self.sim)
self.light_modder = LightModder(self.sim)
if self.FLAGS['domrand']:
self.tex_modder.whiten_materials()
self.goal = {}
# mean camera position defined relative to table center
if FLAGS['baxter']:
self.CAMERA_POS_MEAN = np.array([-0.64, 0.025, 1.0])
else:
self.CAMERA_POS_MEAN = np.array([-0.65, 0.0, 0.68])
# establish values for discretizing
self.X = self.FLAGS['ACTS']['x']
self.Y = self.FLAGS['ACTS']['y']
self.YAW = self.FLAGS['ACTS']['yaw']
self.DIST = self.FLAGS['ACTS']['dist']
self.action_space = self.FLAGS['action_space']
self.observation_space = self.FLAGS['observation_space']
if self.FLAGS['goal_conditioned']:
self.goal['array'] = np.zeros(self.FLAGS['state_shape'])
if self.FLAGS['use_image']:
self.goal['image'] = np.zeros(self.FLAGS['image_shape'])
self.paddle_bid = self.name2bid('paddle')
self.paddle_gid = self.name2gid('paddle')
self.N = int(
0.5 /
self.dt) # N should correspond to about 0.5 seconds of timestep
self.consecutive_fences = 0
self.escape_count = 0
self._set_invisiwall()
self._reset_sim()