47 lines
1.9 KiB
Python
47 lines
1.9 KiB
Python
import gym
|
|
from gym.wrappers.time_limit import TimeLimit
|
|
import numpy as np
|
|
from intersim.envs.intersimple import NRasterizedRouteRandomAgentLocation, RandomLocation, RandomAgent, RewardVisualization, Reward, \
|
|
ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedObservation, \
|
|
NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple
|
|
|
|
class RasterizedSpeed:
|
|
|
|
def __init__(self, max_speed=12, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
channels, height, width = self.observation_space.shape
|
|
self.observation_space = gym.spaces.Box(
|
|
low=0,
|
|
high=255,
|
|
shape=(channels+1, height, width),
|
|
dtype=np.uint8
|
|
)
|
|
self._max_speed = max_speed
|
|
|
|
def _simple_obs(self, intersim_obs, intersim_info):
|
|
img = super()._simple_obs(intersim_obs, intersim_info)
|
|
|
|
ego_speed = intersim_obs['state'][self._agent, 2]
|
|
scaled_speed = (255 * ego_speed) // self._max_speed
|
|
speed_layer = scaled_speed * np.ones_like(img[:1], dtype=np.uint8)
|
|
speed_layer = speed_layer.clamp(0, 255)
|
|
|
|
obs = np.concatenate((img, speed_layer), axis=0)
|
|
return obs
|
|
|
|
class NRasterizedRouteSpeedRandomAgentLocation(RandomLocation, RandomAgent, RewardVisualization,
|
|
Reward, ImageObservationAnimation, RasterizedRoute, NObservations, RasterizedSpeed, RasterizedObservation,
|
|
NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple):
|
|
pass
|
|
|
|
class TransparentTimeLimit(TimeLimit):
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self.env, name)
|
|
|
|
def close(self, *args, **kwargs):
|
|
return self.env.close(*args, **kwargs)
|
|
|
|
def TLNRasterizedRouteRandomAgentLocation(max_episode_steps, *args, **kwargs):
|
|
return TransparentTimeLimit(NRasterizedRouteRandomAgentLocation(*args, **kwargs), max_episode_steps=max_episode_steps)
|