def test_viewer():
model = load_model_from_path("mujoco_py/tests/test.xml")
sim = MjSim(model)
viewer = MjViewer(sim)
for _ in range(100):
sim.step()
viewer.render()
def render(self, mode='human', close=False):
if close:
# TODO: actually close the inspection viewer
return
assert self.sim is not None, \
"Please reset environment before render()."
if mode == 'human':
# Use a nicely-interactive version of the mujoco viewer
if self.viewer is None:
# Inline import since this is only relevant on platforms
# which have GLFW support.
from mujoco_py.mjviewer import MjViewer # noqa
self.viewer = MjViewer(self.sim)
self.viewer.render()
elif mode == 'rgb_array':
return self.sim.render(500, 500)
else:
raise ValueError("Unsupported mode %s" % mode)
def mujoco_viewer(self):
"""
Get a nicely-interactive version of the mujoco viewer
"""
if self._mujoco_viewer is None:
# Inline import since this is only relevant on platforms
# which have GLFW support.
from mujoco_py.mjviewer import MjViewer # noqa
self._mujoco_viewer = MjViewer(self.sim)
return self._mujoco_viewer
def render(self, width=None, height=None, *, camera_name=None, depth=False,
mode='offscreen', device_id=-1):
"""
Renders view from a camera and returns image as an `numpy.ndarray`.
Args:
- width (int): desired image width.
- height (int): desired image height.
- camera_name (str): name of camera in model. If None, the free
camera will be used.
- depth (bool): if True, also return depth buffer
- device (int): device to use for rendering (only for GPU-backed
rendering).
Returns:
- rgb (uint8 array): image buffer from camera
- depth (float array): depth buffer from camera (only returned
if depth=True)
"""
if camera_name is None:
camera_id = None
else:
camera_id = self.model.camera_name2id(camera_name)
if mode == 'offscreen':
with _MjSim_render_lock:
if self._render_context_offscreen is None:
render_context = MjRenderContextOffscreen(
self, device_id=device_id)
else:
render_context = self._render_context_offscreen
render_context.render(
width=width, height=height, camera_id=camera_id)
return render_context.read_pixels(
width, height, depth=depth)
elif mode == 'window':
if self._render_context_window is None:
from mujoco_py.mjviewer import MjViewer
render_context = MjViewer(self)
else:
render_context = self._render_context_window
render_context.render()
else:
raise ValueError("Mode must be either 'window' or 'offscreen'.")
class Env(gym.Env):
metadata = {
'render.modes': ['human', 'rgb_array'],
}
def __init__(self,
get_sim,
get_obs=flatten_get_obs,
get_reward=zero_get_reward,
get_info=empty_get_info,
get_diverged=false_get_diverged,
set_action=None,
action_space=None,
horizon=100,
start_seed=None,
deterministic_mode=False):
"""
Env is a Gym environment subclass tuned for robotics learning
research.
Args:
- get_sim (callable): a callable that returns an MjSim.
- get_obs (callable): callable with an MjSim object as the sole
argument and should return observations.
- set_action (callable): callable which takes an MjSim object and
updates its data and buffer directly.
- get_reward (callable): callable which takes an MjSim object and
returns a scalar reward.
- get_info (callable): callable which takes an MjSim object and
returns info (dictionary).
- get_diverged (callable): callable which takes an MjSim object
and returns a (bool, float) tuple. First value is True if
simulator diverged and second value is the reward at divergence.
- action_space: a space of allowed actions or a two-tuple of a ranges
if number of actions is unknown until the simulation is instantiated
- horizon (int): horizon of environment (i.e. max number of steps).
- start_seed (int or string): seed for random state generator (None for random seed).
Strings will be hashed. A non-None value implies deterministic_mode=True.
This argument allows us to run a deterministic series of goals/randomizations
for a given policy. Then applying the same seed to another policy will allow the
comparison of results more accurately. The reason a string is allowed is so
that we can more easily find and share seeds that are farther from 0,
which is the default starting point for deterministic_mode, and thus have
more likelihood of getting a performant sequence of goals.
"""
if (horizon is not None) and not isinstance(horizon, int):
raise TypeError('horizon must be an int')
self.get_sim = enforce_is_callable(
get_sim,
('get_sim should be callable and should return an MjSim object'))
self.get_obs = enforce_is_callable(
get_obs,
('get_obs should be callable with an MjSim object as the sole '
'argument and should return observations'))
self.set_action = enforce_is_callable(
set_action,
('set_action should be a callable which takes an MjSim object and '
'updates its data and buffer directly'))
self.get_reward = enforce_is_callable(
get_reward,
('get_reward should be a callable which takes an MjSim object and '
'returns a scalar reward'))
self.get_info = enforce_is_callable(
get_info,
('get_info should be a callable which takes an MjSim object and '
'returns a dictionary'))
self.get_diverged = enforce_is_callable(get_diverged, (
'get_diverged should be a callable which takes an MjSim object '
'and returns a (bool, float) tuple. First value is whether '
'simulator is diverged (or done) and second value is the reward at '
'that time.'))
self.sim = None
self.horizon = horizon
self.t = None
self.deterministic_mode = deterministic_mode
# Numpy Random State
if isinstance(start_seed, str):
start_seed = int(
hashlib.sha1(start_seed.encode()).hexdigest(), 16) % (2**32)
self.deterministic_mode = True
elif isinstance(start_seed, int):
self.deterministic_mode = True
else:
start_seed = 0 if self.deterministic_mode else np.random.randint(
2**32, dtype='int64')
self._random_state = np.random.RandomState(start_seed)
# Seed that will be used on next _reset()
self._next_seed = start_seed
# Seed that was used in last _reset()
self._current_seed = None
# For rendering
self.viewer = None
self.rgbviewer = None
# These are required by Gym
self._action_space = action_space
self._observation_space = None
self._spec = Spec(max_episode_steps=horizon, timestep_limit=horizon)
self._name = None
# This is to mitigate issues with old/new envs
@property
def unwrapped(self):
return self
@property
def name(self):
if self._name is None:
name = str(inspect.getfile(self.get_sim))
if name.endswith(".py"):
name = name[:-3]
self._name = name
return self._name
def set_state(self, state, call_forward=True):
"""
Sets the state of the enviroment to the given value. It does not
set time.
Warning: This only sets the MuJoCo state by setting qpos/qvel
(and the user-defined state "udd_state"). It doesn't set
the state of objects which don't have joints.
Args:
- state (MjSimState): desired state.
- call_forward (bool): if True, forward simulation after setting
state.
"""
if not isinstance(state, MjSimState):
raise TypeError("state must be an MjSimState")
if self.sim is None:
raise EmptyEnvException(
"You must call reset() or reset_to_state() before setting the "
"state the first time")
# Call forward to write out values in the MuJoCo data.
# Note: if udd_callback is set on the MjSim instance, then the
# user will need to call forward() manually before calling step.
self.sim.set_state(state)
if call_forward:
self.sim.forward()
def get_state(self):
"""
Returns a copy of the current environment state.
Returns:
- state (MjSimState): state of the environment's MjSim object.
"""
if self.sim is None:
raise EmptyEnvException(
"You must call reset() or reset_to_state() before accessing "
"the state the first time")
return self.sim.get_state()
def get_xml(self):
'''
:return: full state of the simulator serialized as XML (won't contain
meshes, textures, and data information).
'''
return self.sim.model.get_xml()
def get_mjb(self):
'''
:return: full state of the simulator serialized as mjb.
'''
return self.sim.model.get_mjb()
def reset_to_state(self, state, call_forward=True):
"""
Reset to given state.
Args:
- state (MjSimState): desired state.
"""
if not isinstance(state, MjSimState):
raise TypeError(
"You must reset to an explicit state (MjSimState).")
if self.sim is None:
if self._current_seed is None:
self._update_seed()
self.sim = self.get_sim(self._current_seed)
else:
# Ensure environment state not captured in MuJoCo's qpos/qvel
# is reset to the state defined by the model.
self.sim.reset()
self.set_state(state, call_forward=call_forward)
self.t = 0
return self._reset_sim_and_spaces()
def _update_seed(self, force_seed=None):
if force_seed is not None:
self._next_seed = force_seed
self._current_seed = self._next_seed
assert self._current_seed is not None
# if in deterministic mode, then simply increment seed, otherwise randomize
if self.deterministic_mode:
self._next_seed = self._next_seed + 1
else:
self._next_seed = np.random.randint(2**32, dtype='int64')
# immediately update the seed in the random state object
self._random_state.seed(self._current_seed)
@property
def current_seed(self):
# Note: this is a property rather than just instance variable
# for legacy and backwards compatibility reasons.
return self._current_seed
def _reset_sim_and_spaces(self):
obs = self.get_obs(self.sim)
# Mocaps are defined by 3-dim position and 4-dim quaternion
if isinstance(self._action_space, tuple):
assert len(self._action_space) == 2
self._action_space = Box(
self._action_space[0], self._action_space[1],
(self.sim.model.nmocap * 7 + self.sim.model.nu, ), np.float32)
elif self._action_space is None:
self._action_space = Box(
-np.inf, np.inf,
(self.sim.model.nmocap * 7 + self.sim.model.nu, ), np.float32)
self._action_space.flatten_dim = np.prod(self._action_space.shape)
self._observation_space = gym_space_from_arrays(obs)
if self.viewer is not None:
self.viewer.update_sim(self.sim)
return obs
#
# Custom pickling
#
def __getstate__(self):
excluded_attrs = frozenset(("sim", "viewer", "_monitor"))
attr_values = {
k: v
for k, v in self.__dict__.items() if k not in excluded_attrs
}
if self.sim is not None:
attr_values['sim_state'] = self.get_state()
return attr_values
def __setstate__(self, attr_values):
for k, v in attr_values.items():
if k != 'sim_state':
self.__dict__[k] = v
self.sim = None
self.viewer = None
if 'sim_state' in attr_values:
if self.sim is None:
assert self._current_seed is not None
self.sim = self.get_sim(self._current_seed)
self.set_state(attr_values['sim_state'])
self._reset_sim_and_spaces()
return self
def logs(self):
logs = []
if hasattr(self.env, 'logs'):
logs += self.env.logs()
return logs
#
# GYM REQUIREMENTS: these are methods required to be compatible with Gym
#
@property
def action_space(self):
if self._action_space is None:
raise EmptyEnvException(
"You have to reset environment before accessing action_space.")
return self._action_space
@property
def observation_space(self):
if self._observation_space is None:
raise EmptyEnvException(
"You have to reset environment before accessing "
"observation_space.")
return self._observation_space
def reset(self, force_seed=None):
self._update_seed(force_seed=force_seed)
# get sim with current seed
self.sim = self.get_sim(self._current_seed)
# init sim
self.sim.forward()
self.t = 0
self.sim.data.time = 0.0
return self._reset_sim_and_spaces()
def seed(self, seed=None):
"""
Use `env.seed(some_seed)` to set the seed that'll be used in
`env.reset()`. More specifically, this is the seed that will
be passed into `env.get_sim` during `env.reset()`. The seed
will then be incremented in consequent calls to `env.reset()`.
For example:
env.seed(0)
env.reset() -> gives seed(0) world
env.reset() -> gives seed(1) world
...
env.seed(0)
env.reset() -> gives seed(0) world
"""
if isinstance(seed, list):
# Support list of seeds as required by Gym.
assert len(seed) == 1, "Only a single seed supported."
self._next_seed = seed[0]
elif isinstance(seed, int):
self._next_seed = seed
elif seed is not None:
# If seed is None, we just return current seed.
raise ValueError("Seed must be an integer.")
# Return list of seeds to conform to Gym specs
return [self._next_seed]
def step(self, action):
action = action['action_movement'].flatten()
action = np.asarray(action)
action = np.minimum(action, self.action_space.high)
action = np.maximum(action, self.action_space.low)
assert self.action_space.contains(action), (
'Action should be in action_space:\nSPACE=%s\nACTION=%s' %
(self.action_space, action))
self.set_action(self.sim, action)
self.sim.step()
# Need to call forward() so that sites etc are updated,
# since they're used in the reward computations.
self.sim.forward()
self.t += 1
reward = self.get_reward(self.sim)
if not isinstance(reward, float):
raise TypeError("The return value of get_reward must be a float")
obs = self.get_obs(self.sim)
diverged, divergence_reward = self.get_diverged(self.sim)
if not isinstance(diverged, bool):
raise TypeError(
"The first return value of get_diverged must be boolean")
if not isinstance(divergence_reward, float):
raise TypeError(
"The second return value of get_diverged must be float")
if diverged:
done = True
if divergence_reward is not None:
reward = divergence_reward
elif self.horizon is not None:
done = (self.t >= self.horizon)
else:
done = False
info = self.get_info(self.sim)
info["diverged"] = divergence_reward
# Return value as required by Gym
return obs, reward, done, info
def observe(self):
""" Gets a new observation from the environment. """
self.sim.forward()
return self.get_obs(self.sim)
def render(self, mode='human', close=False):
if close:
# TODO: actually close the inspection viewer
return
assert self.sim is not None, \
"Please reset environment before render()."
if mode == 'human':
# Use a nicely-interactive version of the mujoco viewer
if self.viewer is None:
# Inline import since this is only relevant on platforms
# which have GLFW support.
from mujoco_py.mjviewer import MjViewer # noqa
self.viewer = MjViewer(self.sim)
self.viewer.render()
elif mode == 'rgb_array':
#return self.sim.render(500, 500)
if self.rgbviewer is None:
from environment.worldgen.rgbviewer import RgbViewerBasic
self.rgbviewer = RgbViewerBasic(self.sim)
return self.rgbviewer.render()
else:
raise ValueError("Unsupported mode %s" % mode)