envs.py

import logging
import time

import cv2
from gym.spaces.box import Box
import numpy as np
import gym
from gym import spaces
from gym.envs.registration import register

import universe
from universe import vectorized
from universe.wrappers import BlockingReset, GymCoreAction, EpisodeID, Unvectorize, Vectorize, Vision, Logger
from universe import spaces as vnc_spaces
from universe.spaces.vnc_event import keycode

import tensorflow as tf

from adversarial import RandomNoiseWrapper, FGSMNoiseWrapper

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
universe.configure_logging()

# Register custom environments.
register(
    id='CartPolePixels-v0',
    entry_point='pixels.cartpole:CartPolePixelsEnv',
    tags={'wrapper_config.TimeLimit.max_episode_steps': 200},
    reward_threshold=195.0,
)

register(
    id='CartPolePixels-IncreasedGravity-v0',
    entry_point='pixels.cartpole:CartPolePixelsEnv',
    tags={'wrapper_config.TimeLimit.max_episode_steps': 200},
    kwargs={
        'gravity': 20.0,
    },
    reward_threshold=195.0,
)

register(
    id='CartPolePixels-DecreasedGravity-v0',
    entry_point='pixels.cartpole:CartPolePixelsEnv',
    tags={'wrapper_config.TimeLimit.max_episode_steps': 200},
    kwargs={
        'gravity': 5.0,
    },
    reward_threshold=195.0,
)

register(
    id='CartPolePixels-IncreasedCartMass-v0',
    entry_point='pixels.cartpole:CartPolePixelsEnv',
    tags={'wrapper_config.TimeLimit.max_episode_steps': 200},
    kwargs={
        'mass_cart': 2.0,
    },
    reward_threshold=195.0,
)

register(
    id='CartPolePixels-IncreasedPoleMass-v0',
    entry_point='pixels.cartpole:CartPolePixelsEnv',
    tags={'wrapper_config.TimeLimit.max_episode_steps': 200},
    kwargs={
        'mass_pole': 0.3,
    },
    reward_threshold=195.0,
)

register(
    id='CarRacingPixels-v0',
    entry_point='pixels.car_racing:CarRacingPixelsEnv',
    tags={'wrapper_config.TimeLimit.max_episode_steps': 1000},
    reward_threshold=900,
)


def create_env(env_id, client_id, remotes, **kwargs):
    spec = gym.spec(env_id)

    if spec.tags.get('flashgames', False):
        return create_flash_env(env_id, client_id, remotes, **kwargs)
    elif spec.tags.get('atari', False) and spec.tags.get('vnc', False):
        return create_vncatari_env(env_id, client_id, remotes, **kwargs)
    else:
        # Assume atari.
        assert "." not in env_id  # universe environments have dots in names.

        if env_id.startswith('CartPolePixels-') or env_id.startswith('CarRacingPixels-'):
            return create_pixels_env(env_id, **kwargs)
        else:
            return create_atari_env(env_id, **kwargs)


def create_adversarial_env(env, adversarial_mode=None, adversarial_epsilon=0.01, refs=None, **kwargs):
    """
    Create an adversarial environment.
    """

    setup_functions = []

    def adversary_setup(*args, **kwargs):
        for function in setup_functions:
            function(*args, **kwargs)

    if adversarial_mode == 'random':
        env = RandomNoiseWrapper(env, intensity=adversarial_epsilon)
    elif adversarial_mode == 'fgsm':
        env = FGSMNoiseWrapper(env, intensity=adversarial_epsilon)
        setup_functions.append(env.setup)
    elif adversarial_mode == 'fgsm_skip5':
        env = FGSMNoiseWrapper(env, intensity=adversarial_epsilon, skip=5)
        setup_functions.append(env.setup)
    elif adversarial_mode == 'fgsm_skip10':
        env = FGSMNoiseWrapper(env, intensity=adversarial_epsilon, skip=10)
        setup_functions.append(env.setup)
    elif adversarial_mode == 'fgsm_reuse5':
        env = FGSMNoiseWrapper(env, intensity=adversarial_epsilon, skip=5, reuse=True)
        setup_functions.append(env.setup)
    elif adversarial_mode == 'fgsm_reuse10':
        env = FGSMNoiseWrapper(env, intensity=adversarial_epsilon, skip=10, reuse=True)
        setup_functions.append(env.setup)
    elif adversarial_mode == 'fgsm_vf':
        env = FGSMNoiseWrapper(env, intensity=adversarial_epsilon, vf=True)
        setup_functions.append(env.setup)

    refs['adversary_setup'] = adversary_setup

    return env


def create_flash_env(env_id, client_id, remotes, **kwargs):
    env = gym.make(env_id)
    env = Vision(env)
    env = Logger(env)
    env = BlockingReset(env)

    reg = universe.runtime_spec('flashgames').server_registry
    height = reg[env_id]["height"]
    width = reg[env_id]["width"]
    env = CropScreen(env, height, width, 84, 18)
    env = FlashRescale(env)

    keys = ['left', 'right', 'up', 'down', 'x']
    if env_id == 'flashgames.NeonRace-v0':
        # Better key space for this game.
        keys = ['left', 'right', 'up', 'left up', 'right up', 'down', 'up x']
    logger.info('create_flash_env(%s): keys=%s', env_id, keys)

    env = DiscreteToFixedKeysVNCActions(env, keys)
    env = EpisodeID(env)
    env = DiagnosticsInfo(env)
    env = create_adversarial_env(env, **kwargs)
    env = Unvectorize(env)
    env.configure(fps=5.0, remotes=remotes, start_timeout=15 * 60, client_id=client_id,
                  vnc_driver='go', vnc_kwargs={
                    'encoding': 'tight', 'compress_level': 0,
                    'fine_quality_level': 50, 'subsample_level': 3})
    return env


def create_vncatari_env(env_id, client_id, remotes, **kwargs):
    env = gym.make(env_id)
    env = Vision(env)
    env = Logger(env)
    env = BlockingReset(env)
    env = GymCoreAction(env)
    env = AtariRescale42x42(env)
    env = EpisodeID(env)
    env = DiagnosticsInfo(env)
    env = create_adversarial_env(env, **kwargs)
    env = Unvectorize(env)

    logger.info('Connecting to remotes: %s', remotes)
    fps = env.metadata['video.frames_per_second']
    env.configure(remotes=remotes, start_timeout=15 * 60, fps=fps, client_id=client_id)
    return env


def create_atari_env(env_id, **kwargs):
    env = gym.make(env_id)

    # Seed.
    env.seed(0)
    np.random.seed(0)
    tf.set_random_seed(0)

    env = Vectorize(env)
    env = AtariRescale42x42(env)
    env = create_adversarial_env(env, **kwargs)
    env = DiagnosticsInfo(env)
    env = Unvectorize(env)
    return env


def create_pixels_env(env_id, **kwargs):
    env = gym.make(env_id)
    env = Vectorize(env)
    env = create_adversarial_env(env, **kwargs)
    env = DiagnosticsInfo(env)
    env = Unvectorize(env)
    return env


def DiagnosticsInfo(env, *args, **kwargs):
    return vectorized.VectorizeFilter(env, DiagnosticsInfoI, *args, **kwargs)


class DiagnosticsInfoI(vectorized.Filter):
    def __init__(self, log_interval=503):
        super(DiagnosticsInfoI, self).__init__()

        self._episode_time = time.time()
        self._last_time = time.time()
        self._local_t = 0
        self._log_interval = log_interval
        self._episode_reward = 0
        self._episode_length = 0
        self._all_rewards = []
        self._num_vnc_updates = 0
        self._last_episode_id = -1
        self._adversary = {}

    def _after_reset(self, observation):
        logger.info('Resetting environment')
        self._episode_reward = 0
        self._episode_length = 0
        self._all_rewards = []
        self._adversary = {}
        return observation

    def _after_step(self, observation, reward, done, info):
        to_log = {}
        if self._episode_length == 0:
            self._episode_time = time.time()

        self._local_t += 1
        if info.get("stats.vnc.updates.n") is not None:
            self._num_vnc_updates += info.get("stats.vnc.updates.n")

        if self._local_t % self._log_interval == 0:
            cur_time = time.time()
            elapsed = cur_time - self._last_time
            fps = self._log_interval / elapsed
            self._last_time = cur_time
            cur_episode_id = info.get('vectorized.episode_id', 0)
            to_log["diagnostics/fps"] = fps
            if self._last_episode_id == cur_episode_id:
                to_log["diagnostics/fps_within_episode"] = fps
            self._last_episode_id = cur_episode_id
            if info.get("stats.gauges.diagnostics.lag.action") is not None:
                to_log["diagnostics/action_lag_lb"] = info["stats.gauges.diagnostics.lag.action"][0]
                to_log["diagnostics/action_lag_ub"] = info["stats.gauges.diagnostics.lag.action"][1]
            if info.get("reward.count") is not None:
                to_log["diagnostics/reward_count"] = info["reward.count"]
            if info.get("stats.gauges.diagnostics.clock_skew") is not None:
                to_log["diagnostics/clock_skew_lb"] = info["stats.gauges.diagnostics.clock_skew"][0]
                to_log["diagnostics/clock_skew_ub"] = info["stats.gauges.diagnostics.clock_skew"][1]
            if info.get("stats.gauges.diagnostics.lag.observation") is not None:
                to_log["diagnostics/observation_lag_lb"] = info["stats.gauges.diagnostics.lag.observation"][0]
                to_log["diagnostics/observation_lag_ub"] = info["stats.gauges.diagnostics.lag.observation"][1]

            if info.get("stats.vnc.updates.n") is not None:
                to_log["diagnostics/vnc_updates_n"] = info["stats.vnc.updates.n"]
                to_log["diagnostics/vnc_updates_n_ps"] = self._num_vnc_updates / elapsed
                self._num_vnc_updates = 0
            if info.get("stats.vnc.updates.bytes") is not None:
                to_log["diagnostics/vnc_updates_bytes"] = info["stats.vnc.updates.bytes"]
            if info.get("stats.vnc.updates.pixels") is not None:
                to_log["diagnostics/vnc_updates_pixels"] = info["stats.vnc.updates.pixels"]
            if info.get("stats.vnc.updates.rectangles") is not None:
                to_log["diagnostics/vnc_updates_rectangles"] = info["stats.vnc.updates.rectangles"]
            if info.get("env_status.state_id") is not None:
                to_log["diagnostics/env_state_id"] = info["env_status.state_id"]

        if reward is not None:
            self._episode_reward += reward
            if observation is not None:
                self._episode_length += 1
            self._all_rewards.append(reward)

        for key, value in info.items():
            if key.startswith('adversary/'):
                self._adversary.setdefault(key, []).append(value)

        if done:
            logger.info('Episode terminating: episode_reward={} episode_length={}'.format(
                self._episode_reward, self._episode_length
            ))
            total_time = time.time() - self._episode_time
            to_log["global/episode_reward"] = self._episode_reward
            to_log["global/episode_length"] = self._episode_length
            to_log["global/episode_time"] = total_time
            to_log["global/reward_per_time"] = self._episode_reward / total_time

            for key, value in self._adversary.items():
                to_log[key] = np.average(value)

            self._episode_reward = 0
            self._episode_length = 0
            self._all_rewards = []
            self._adversary = {}

        return observation, reward, done, to_log


def _process_frame42(frame):
    frame = frame[34:34+160, :160]
    # Resize by half, then down to 42x42 (essentially mipmapping). If
    # we resize directly we lose pixels that, when mapped to 42x42,
    # aren't close enough to the pixel boundary.
    frame = cv2.resize(frame, (80, 80))
    frame = cv2.resize(frame, (42, 42))
    frame = frame.mean(2)
    frame = frame.astype(np.float32)
    frame *= (1.0 / 255.0)
    frame = np.reshape(frame, [42, 42, 1])
    return frame


class AtariRescale42x42(vectorized.ObservationWrapper):
    def __init__(self, env=None):
        super(AtariRescale42x42, self).__init__(env)
        self.observation_space = Box(0.0, 1.0, [42, 42, 1])

    def _observation(self, observation_n):
        return [_process_frame42(observation) for observation in observation_n]


class FixedKeyState(object):
    def __init__(self, keys):
        self._keys = [keycode(key) for key in keys]
        self._down_keysyms = set()

    def apply_vnc_actions(self, vnc_actions):
        for event in vnc_actions:
            if isinstance(event, vnc_spaces.KeyEvent):
                if event.down:
                    self._down_keysyms.add(event.key)
                else:
                    self._down_keysyms.discard(event.key)

    def to_index(self):
        action_n = 0
        for key in self._down_keysyms:
            if key in self._keys:
                # If multiple keys are pressed, just use the first one
                action_n = self._keys.index(key) + 1
                break
        return action_n


class DiscreteToFixedKeysVNCActions(vectorized.ActionWrapper):
    """
    Define a fixed action space. Action 0 is all keys up. Each element of keys can
    be a single key or a space-separated list of keys

    For example,
       e=DiscreteToFixedKeysVNCActions(e, ['left', 'right'])
    will have 3 actions: [none, left, right]

    You can define a state with more than one key down by separating with spaces. For example,
       e=DiscreteToFixedKeysVNCActions(e, ['left', 'right', 'space', 'left space', 'right space'])
    will have 6 actions: [none, left, right, space, left space, right space]
    """
    def __init__(self, env, keys):
        super(DiscreteToFixedKeysVNCActions, self).__init__(env)

        self._keys = keys
        self._generate_actions()
        self.action_space = spaces.Discrete(len(self._actions))

    def _generate_actions(self):
        self._actions = []
        uniq_keys = set()
        for key in self._keys:
            for cur_key in key.split(' '):
                uniq_keys.add(cur_key)

        for key in [''] + self._keys:
            split_keys = key.split(' ')
            cur_action = []
            for cur_key in uniq_keys:
                cur_action.append(vnc_spaces.KeyEvent.by_name(cur_key, down=(cur_key in split_keys)))
            self._actions.append(cur_action)
        self.key_state = FixedKeyState(uniq_keys)

    def _action(self, action_n):
        # Each action might be a length-1 np.array. Cast to int to
        # avoid warnings.
        return [self._actions[int(action)] for action in action_n]


class CropScreen(vectorized.ObservationWrapper):
    """Crops out a [height]x[width] area starting from (top,left) """
    def __init__(self, env, height, width, top=0, left=0):
        super(CropScreen, self).__init__(env)
        self.height = height
        self.width = width
        self.top = top
        self.left = left
        self.observation_space = Box(0, 255, shape=(height, width, 3))

    def _observation(self, observation_n):
        return [ob[self.top:self.top+self.height, self.left:self.left+self.width, :] if ob is not None else None
                for ob in observation_n]


def _process_frame_flash(frame):
    frame = cv2.resize(frame, (200, 128))
    frame = frame.mean(2).astype(np.float32)
    frame *= (1.0 / 255.0)
    frame = np.reshape(frame, [128, 200, 1])
    return frame


class FlashRescale(vectorized.ObservationWrapper):
    def __init__(self, env=None):
        super(FlashRescale, self).__init__(env)
        self.observation_space = Box(0.0, 1.0, [128, 200, 1])

    def _observation(self, observation_n):
        return [_process_frame_flash(observation) for observation in observation_n]