Spaces:
Running on Zero
Running on Zero
File size: 3,100 Bytes
23a59ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | import gymnasium as gym
import numpy as np
from envs.tasks.dmcontrol import cartpole, cheetah, walker, hopper, reacher, ball_in_cup, pendulum, fish, giraffe, spinner, jumper, finger
from dm_control import suite
suite._DOMAINS['giraffe'] = giraffe
suite._DOMAINS['spinner'] = spinner
suite._DOMAINS['jumper'] = jumper
suite.ALL_TASKS = suite.ALL_TASKS + suite._get_tasks('custom')
suite.TASKS_BY_DOMAIN = suite._get_tasks_by_domain(suite.ALL_TASKS)
from dm_control.suite.wrappers import action_scale
from envs.wrappers.timeout import Timeout
from envs.wrappers.pixels import Pixels
def get_obs_shape(env):
obs_shp = []
for v in env.observation_spec().values():
try:
shp = np.prod(v.shape)
except:
shp = 1
obs_shp.append(shp)
return (int(np.sum(obs_shp)),)
class DMControlWrapper:
def __init__(self, env, domain):
self.env = env
self.camera_id = 2 if domain == 'quadruped' else 0
obs_shape = get_obs_shape(env)
action_shape = env.action_spec().shape
self.observation_space = gym.spaces.Box(
low=np.full(obs_shape, -np.inf, dtype=np.float32),
high=np.full(obs_shape, np.inf, dtype=np.float32),
dtype=np.float32)
self.action_space = gym.spaces.Box(
low=np.full(action_shape, env.action_spec().minimum),
high=np.full(action_shape, env.action_spec().maximum),
dtype=env.action_spec().dtype)
self.action_spec_dtype = env.action_spec().dtype
self._cumulative_reward = 0
@property
def unwrapped(self):
return self.env
@property
def metadata(self):
return None
@property
def info(self):
return {
'terminated': False,
'truncated': False,
'success': float('nan'),
'score': self._cumulative_reward/1000,
}
def _obs_to_array(self, obs):
return np.concatenate([v.flatten() for v in obs.values()], dtype=np.float32)
def reset(self):
self._cumulative_reward = 0
return self._obs_to_array(self.env.reset().observation), self.info
def step(self, action):
reward = 0
action = action.astype(self.action_spec_dtype)
for _ in range(2):
step = self.env.step(action)
reward += step.reward
self._cumulative_reward += reward
return self._obs_to_array(step.observation), reward, False, False, self.info
def render(self, width=224, height=224, camera_id=None):
return self.env.physics.render(height, width, camera_id or self.camera_id)
def close(self):
self.env.close()
def make_env(cfg):
"""
Make DMControl environment.
Adapted from https://github.com/facebookresearch/drqv2
"""
domain, task = cfg.task.replace('-', '_').split('_', 1)
domain = dict(cup='ball_in_cup', pointmass='point_mass').get(domain, domain)
if (domain, task) not in suite.ALL_TASKS:
raise ValueError('Unknown task:', task)
assert cfg.obs in {'state', 'rgb'}, 'This task only supports state and rgb observations.'
env = suite.load(domain,
task,
task_kwargs={'random': cfg.seed},
visualize_reward=False)
env = action_scale.Wrapper(env, minimum=-1., maximum=1.)
env = DMControlWrapper(env, domain)
if cfg.obs == 'rgb':
env = Pixels(env, cfg)
env = Timeout(env, max_episode_steps=500)
return env
|