112 lines
3.9 KiB
Python
112 lines
3.9 KiB
Python
import gym
|
|
import numpy as np
|
|
from src.gail2.wrappers import Wrapper, Setobs, TransformObservation
|
|
from intersim.envs import IntersimpleLidarFlatIncrementingAgent
|
|
|
|
obs_min = np.array([
|
|
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
]).reshape(-1)
|
|
|
|
obs_max = np.array([
|
|
[1000, 1000, 20, np.pi, 1e-1, 0.],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
]).reshape(-1)
|
|
|
|
def NormalizedOptionsEvalEnv(**kwargs):
|
|
return OptionsEnv(Setobs(
|
|
TransformObservation(IntersimpleLidarFlatIncrementingAgent(
|
|
n_rays=5,
|
|
**kwargs,
|
|
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
|
|
), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5)])
|
|
|
|
def NormalizedContinuousEvalEnv(**kwargs):
|
|
return Setobs(
|
|
TransformObservation(IntersimpleLidarFlatIncrementingAgent(
|
|
n_rays=5,
|
|
**kwargs,
|
|
), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
|
|
)
|
|
|
|
class OptionsEnv(Wrapper):
|
|
|
|
def __init__(self, env, options):
|
|
super().__init__(env)
|
|
self.ll_action_space = env.action_space
|
|
self.options = options
|
|
self.action_space = gym.spaces.Discrete(len(options))
|
|
self.max_plan_length = max(t for _, t in options)
|
|
|
|
def plan(self, option):
|
|
target_v, t = option
|
|
current_v = self.env._env.state[self.env._agent, 1].item()
|
|
dt = self.env._env._dt
|
|
a = (target_v - current_v) / (t * dt)
|
|
a = self.env._normalize(a)
|
|
a = a * np.ones((t,))
|
|
a += 0.01 * np.random.randn(*a.shape)
|
|
a = np.clip(a, self.ll_action_space.low, self.ll_action_space.high)
|
|
return a
|
|
|
|
def execute_plan(self, obs, option, render_mode=None):
|
|
observations = np.zeros((self.max_plan_length + 1, *self.env.observation_space.shape))
|
|
actions = np.zeros((self.max_plan_length + 1, *self.ll_action_space.shape))
|
|
rewards = np.zeros((self.max_plan_length + 1,))
|
|
env_done = np.ones((self.max_plan_length + 1,), dtype=bool)
|
|
plan_done = np.ones((self.max_plan_length + 1,), dtype=bool)
|
|
infos = []
|
|
|
|
observations[0] = obs
|
|
env_done[0] = False
|
|
for k, u in enumerate(self.plan(option)):
|
|
plan_done[k] = False
|
|
o, r, d, i = super().step(u)
|
|
actions[k] = u
|
|
rewards[k] = r
|
|
env_done[k] = d
|
|
infos.append(i)
|
|
observations[k+1] = o
|
|
|
|
if render_mode is not None:
|
|
self.env.render(render_mode)
|
|
|
|
if d:
|
|
break
|
|
|
|
n_steps = k + 1
|
|
return observations, actions, rewards, env_done, plan_done, infos, n_steps
|
|
|
|
def step(self, action, render_mode=None):
|
|
a = int(action)
|
|
assert a == action
|
|
ll_obs, ll_actions, ll_rewards, ll_env_done, ll_plan_done, ll_infos, ll_steps = self.execute_plan(self.last_obs, self.options[a], render_mode)
|
|
hl_obs = ll_obs[ll_steps]
|
|
hl_reward = (ll_rewards * ~ll_plan_done).sum().item()
|
|
hl_done = ll_env_done[ll_steps-1].item()
|
|
hl_infos = {
|
|
'll': {
|
|
'observations': ll_obs,
|
|
'actions': ll_actions,
|
|
'rewards': ll_rewards,
|
|
'env_done': ll_env_done,
|
|
'plan_done': ll_plan_done,
|
|
'infos': ll_infos,
|
|
'steps': ll_steps,
|
|
}
|
|
}
|
|
self.last_obs = hl_obs
|
|
return hl_obs, hl_reward, hl_done, hl_infos
|
|
|
|
def reset(self, *args, **kwargs):
|
|
self.last_obs = super().reset(*args, **kwargs)
|
|
return self.last_obs
|