145 lines
4.1 KiB
Python
145 lines
4.1 KiB
Python
import datetime
|
|
import gym
|
|
import numpy as np
|
|
import uuid
|
|
|
|
|
|
class TimeLimit:
|
|
def __init__(self, env, duration):
|
|
self._env = env
|
|
self._duration = duration
|
|
self._step = None
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._env, name)
|
|
|
|
def step(self, action):
|
|
assert self._step is not None, "Must reset environment."
|
|
obs, reward, done, info = self._env.step(action)
|
|
self._step += 1
|
|
if self._step >= self._duration:
|
|
done = True
|
|
if "discount" not in info:
|
|
info["discount"] = np.array(1.0).astype(np.float32)
|
|
self._step = None
|
|
return obs, reward, done, info
|
|
|
|
def reset(self):
|
|
self._step = 0
|
|
return self._env.reset()
|
|
|
|
|
|
class NormalizeActions:
|
|
def __init__(self, env):
|
|
self._env = env
|
|
self._mask = np.logical_and(
|
|
np.isfinite(env.action_space.low), np.isfinite(env.action_space.high)
|
|
)
|
|
self._low = np.where(self._mask, env.action_space.low, -1)
|
|
self._high = np.where(self._mask, env.action_space.high, 1)
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._env, name)
|
|
|
|
@property
|
|
def action_space(self):
|
|
low = np.where(self._mask, -np.ones_like(self._low), self._low)
|
|
high = np.where(self._mask, np.ones_like(self._low), self._high)
|
|
return gym.spaces.Box(low, high, dtype=np.float32)
|
|
|
|
def step(self, action):
|
|
original = (action + 1) / 2 * (self._high - self._low) + self._low
|
|
original = np.where(self._mask, original, action)
|
|
return self._env.step(original)
|
|
|
|
|
|
class OneHotAction:
|
|
def __init__(self, env):
|
|
assert isinstance(env.action_space, gym.spaces.Discrete)
|
|
self._env = env
|
|
self._random = np.random.RandomState()
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._env, name)
|
|
|
|
@property
|
|
def action_space(self):
|
|
shape = (self._env.action_space.n,)
|
|
space = gym.spaces.Box(low=0, high=1, shape=shape, dtype=np.float32)
|
|
space.sample = self._sample_action
|
|
space.discrete = True
|
|
return space
|
|
|
|
def step(self, action):
|
|
index = np.argmax(action).astype(int)
|
|
reference = np.zeros_like(action)
|
|
reference[index] = 1
|
|
if not np.allclose(reference, action):
|
|
raise ValueError(f"Invalid one-hot action:\n{action}")
|
|
return self._env.step(index)
|
|
|
|
def reset(self):
|
|
return self._env.reset()
|
|
|
|
def _sample_action(self):
|
|
actions = self._env.action_space.n
|
|
index = self._random.randint(0, actions)
|
|
reference = np.zeros(actions, dtype=np.float32)
|
|
reference[index] = 1.0
|
|
return reference
|
|
|
|
|
|
class RewardObs:
|
|
def __init__(self, env):
|
|
self._env = env
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._env, name)
|
|
|
|
@property
|
|
def observation_space(self):
|
|
spaces = self._env.observation_space.spaces
|
|
if "reward" not in spaces:
|
|
spaces["reward"] = gym.spaces.Box(
|
|
-np.inf, np.inf, shape=(1,), dtype=np.float32
|
|
)
|
|
return gym.spaces.Dict(spaces)
|
|
|
|
def step(self, action):
|
|
obs, reward, done, info = self._env.step(action)
|
|
if "reward" not in obs:
|
|
obs["reward"] = reward
|
|
return obs, reward, done, info
|
|
|
|
def reset(self):
|
|
obs = self._env.reset()
|
|
if "reward" not in obs:
|
|
obs["reward"] = 0.0
|
|
return obs
|
|
|
|
|
|
class SelectAction:
|
|
def __init__(self, env, key):
|
|
self._env = env
|
|
self._key = key
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._env, name)
|
|
|
|
def step(self, action):
|
|
return self._env.step(action[self._key])
|
|
|
|
class UUID:
|
|
def __init__(self, env):
|
|
self._env = env
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%dT%H%M%S")
|
|
self.id = f"{timestamp}-{str(uuid.uuid4().hex)}"
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self._env, name)
|
|
|
|
def reset(self):
|
|
timestamp = datetime.datetime.now().strftime("%Y%m%dT%H%M%S")
|
|
self.id = f"{timestamp}-{str(uuid.uuid4().hex)}"
|
|
return self._env.reset()
|