diff --git a/scratch/etienne/intersimple/gail/envs.py b/scratch/etienne/intersimple/gail/envs.py new file mode 100644 index 0000000..7983312 --- /dev/null +++ b/scratch/etienne/intersimple/gail/envs.py @@ -0,0 +1,34 @@ +import gym +import numpy as np +from intersim.envs.intersimple import 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, RasterizedSpeed, NObservations, RasterizedObservation, + NormalizedActionSpace, ActionVisualization, InteractionSimulatorMarkerViz, ImitationCompat, Intersimple): + pass \ No newline at end of file diff --git a/scratch/etienne/intersimple/gail_options_image_random_location.py b/scratch/etienne/intersimple/gail_options_image_random_location.py index 806a45d..511ecff 100644 --- a/scratch/etienne/intersimple/gail_options_image_random_location.py +++ b/scratch/etienne/intersimple/gail_options_image_random_location.py @@ -5,7 +5,7 @@ sys.path.append('../../../') from src.discriminator import CnnDiscriminatorFlatAction from imitation.algorithms import adversarial import stable_baselines3 -from intersim.envs import NRasterizedRouteSpeedRandomAgentLocation +from gail.envs import NRasterizedRouteSpeedRandomAgentLocation import pickle import imitation.data.rollout as rollout import tempfile