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 _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, 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 test_materials():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.premod.png')
random_state = np.random.RandomState(0)
modder = MaterialModder(sim, random_state=random_state)
modder.set_specularity('g1', 1.0)
modder.set_reflectance('g2', 1.0)
modder.set_shininess('g3', 1.0)
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.props.png')
modder.rand_all('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.rand_all.png')
def test_materials():
model = load_model_from_xml(BASIC_MODEL_XML)
sim = MjSim(model)
sim.forward()
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.premod.png')
random_state = np.random.RandomState(0)
modder = MaterialModder(sim, random_state=random_state)
modder.set_specularity('g1', 1.0)
modder.set_reflectance('g2', 1.0)
modder.set_shininess('g3', 1.0)
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.props.png')
modder.rand_all('g4')
compare_imgs(sim.render(201, 205, camera_name="topcam"),
'test_materials.rand_all.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 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