Spaces:
Running on Zero
Running on Zero
File size: 3,397 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 | import numpy as np
import gymnasium as gym
import torch
from torchvision.transforms import functional as F
from envs.tasks.bipedal_walker import BipedalWalkerFlat, BipedalWalkerUneven, \
BipedalWalkerRugged, BipedalWalkerHills, BipedalWalkerObstacles
from envs.tasks.lunarlander import LunarLander, LunarLanderHover, LunarLanderTakeoff
from envs.wrappers.timeout import Timeout
BOX2D_TASKS = {
'bipedal-walker-flat': BipedalWalkerFlat,
'bipedal-walker-uneven': BipedalWalkerUneven,
'bipedal-walker-rugged': BipedalWalkerRugged,
'bipedal-walker-hills': BipedalWalkerHills,
'bipedal-walker-obstacles': BipedalWalkerObstacles,
'lunarlander-land': LunarLander,
'lunarlander-hover': LunarLanderHover,
'lunarlander-takeoff': LunarLanderTakeoff,
}
class Box2DWrapper(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,
})
self._cumulative_reward = 0
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.startswith('lunarlander'):
info['score'] = np.clip(self._cumulative_reward / 500, 0, 1)
else:
info['score'] = np.clip(self._cumulative_reward / 250, 0, 1)
return info
def get_observation(self, obs):
if self.cfg.obs == 'rgb':
return {'state': obs, 'rgb': self.render().transpose(2, 0, 1)}
return obs
def reset(self):
obs, info = self.env.reset()
self._cumulative_reward = 0
return self.get_observation(obs), self._extract_info(info)
def _safe_action(self, action):
"""
Ensure the action is within the valid range for the environment.
"""
assert isinstance(action, np.ndarray), 'Action must be a numpy array.'
assert not np.isnan(action).any(), 'Action contains NaN values.'
action = np.clip(action, self.env.action_space.low, self.env.action_space.high)
return action
def step(self, action):
reward = 0
action = self._safe_action(action)
for _ in range(2):
obs, _reward, _, truncated, info = self.env.step(action.copy())
if self.cfg.task == 'lunarlander-land':
_reward = np.clip(_reward, -5, 10) / 20
if self.cfg.task.startswith('bipedal-walker'):
_reward = np.clip(_reward, -1, 10)
reward += _reward
if truncated:
break
terminated = False
self._cumulative_reward += reward
info['terminated'] = terminated
info['truncated'] = truncated
return self.get_observation(obs), reward, terminated, truncated, self._extract_info(info)
@property
def unwrapped(self):
return self.env.unwrapped
def render(self, **kwargs):
frame = self.env.render()
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 Box2D environment.
"""
if not cfg.task in BOX2D_TASKS:
raise ValueError('Unknown task:', cfg.task)
env = BOX2D_TASKS[cfg.task](render_mode='rgb_array')
env = Box2DWrapper(env, cfg)
env = Timeout(env, max_episode_steps=250)
return env
|