Spaces:
Running on Zero
Running on Zero
File size: 4,327 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | import numpy as np
import gymnasium as gym
import torch
from torchvision.transforms import functional as F
from envs.wrappers.timeout import Timeout
MUJOCO_TASKS = {
'mujoco-ant': 'Ant-v4',
'mujoco-inverted-pendulum': 'InvertedPendulum-v4',
'mujoco-reacher': 'Reacher-v4',
'mujoco-pusher': 'Pusher-v4',
'mujoco-halfcheetah': 'HalfCheetah-v4',
'mujoco-hopper': 'Hopper-v4',
'mujoco-walker': 'Walker2d-v4',
}
class MuJoCoWrapper(gym.Wrapper):
def __init__(self, env, cfg):
super().__init__(env)
self.env = env
self.cfg = cfg
if cfg.obs == 'rgb':
self.observation_space = gym.spaces.Dict({
'rgb': gym.spaces.Box(
low=0, high=255, shape=(3, self.cfg.render_size, self.cfg.render_size), dtype=np.uint8),
'state': env.observation_space,
})
action_dim = env.action_space.shape[0]
self.action_space = gym.spaces.Box(
low=np.full(action_dim, -1),
high=np.full(action_dim, +1),
dtype=np.float32,
)
self.action_scale = self.env.action_space.high
self.action_repeat = 1 if cfg.task in {
'mujoco-reacher', 'mujoco-pusher', 'mujoco-halfcheetah'
} else 2
self._cumulative_reward = 0
self._terminated = False
def _extract_info(self, info):
info = {
'terminated': info.get('terminated', False),
'truncated': info.get('truncated', False),
'success': float(info.get('success', 0.)),
}
if self.cfg.task == 'mujoco-inverted-pendulum':
# range is [0, 1000], normalize to [0, 1]
info['score'] = np.clip(self._cumulative_reward, 0, 1000) / 1000
elif self.cfg.task == 'mujoco-reacher':
# range is [-50, 0], normalize to [0, 1]
info['score'] = 1 + np.clip(self._cumulative_reward, -50, 0) / 50
elif self.cfg.task == 'mujoco-pusher':
# Pusher-v4 over 100 steps: dense penalty = -dist(obj,goal) - 0.1*dist(arm,obj) - 0.001*||a||^2.
# Cumulative range roughly [-150, 0]; expert ~[-50, -20].
info['score'] = 1 + np.clip(self._cumulative_reward, -150, 0) / 150
elif self.cfg.task == 'mujoco-halfcheetah':
# range is [0, 15000], normalize to [0, 1]
info['score'] = np.clip(self._cumulative_reward, 0, 15000) / 15000
elif self.cfg.task in {'mujoco-ant', 'mujoco-hopper', 'mujoco-walker'}:
# range is [0, 5000], normalize to [0, 1]
info['score'] = np.clip(self._cumulative_reward, 0, 5000) / 5000
else:
raise NotImplementedError(f'Score calculation for {self.cfg.task} not implemented.')
return info
def get_observation(self, obs):
if self.cfg.obs == 'rgb':
return {'state': obs, 'rgb': self.render().transpose(2, 0, 1)}
return obs.astype(np.float32)
def reset(self):
obs, info = self.env.reset()
self._cumulative_reward = 0
self._terminated = False
return self.get_observation(obs), self._extract_info(info)
def step(self, action):
action = action * self.action_scale
reward = 0.
for _ in range(self.action_repeat):
obs, _reward, terminated, truncated, info = self.env.step(action.copy())
if 'pendulum' in self.cfg.task and (terminated or self._terminated):
self._terminated = True
_reward = 0.
elif 'hopper' in self.cfg.task or 'walker' in self.cfg.task:
_reward = max(0, _reward) if self.env.unwrapped.is_healthy else -1
reward += _reward
self._cumulative_reward += reward
info['terminated'] = False
info['truncated'] = truncated
return self.get_observation(obs), reward, False, truncated, self._extract_info(info)
@property
def unwrapped(self):
return self.env.unwrapped
def render(self, **kwargs):
frame = self.env.render().copy()
h, w = self.cfg.render_size, self.cfg.render_size
if frame.shape[0] > h or frame.shape[1] > w:
frame = torch.from_numpy(frame).permute(2, 0, 1)
frame = F.resize(frame, (h, w))
frame = frame.permute(1, 2, 0).numpy()
return frame
def make_env(cfg):
"""
Make MuJoCo environment.
"""
if not cfg.task in MUJOCO_TASKS:
raise ValueError('Unknown task:', cfg.task)
if cfg.task in {'mujoco-ant', 'mujoco-hopper', 'mujoco-walker'}:
env = gym.make(
MUJOCO_TASKS[cfg.task],
terminate_when_unhealthy=False,
render_mode='rgb_array',
)
else:
env = gym.make(
MUJOCO_TASKS[cfg.task],
render_mode='rgb_array',
)
env = MuJoCoWrapper(env, cfg)
env = Timeout(env, max_episode_steps={
'mujoco-reacher': 50,
'mujoco-pusher': 100,
'mujoco-halfcheetah': 1000,
}.get(cfg.task, 500))
return env
|