"""Utilities of visualising an environment.""" from collections import deque from typing import Callable, Dict, List, Optional, Tuple, Union import numpy as np import gym.error from gym import Env, logger from gym.core import ActType, ObsType from gym.error import DependencyNotInstalled from gym.logger import deprecation try: import pygame from pygame import Surface from pygame.event import Event from pygame.locals import VIDEORESIZE except ImportError: raise gym.error.DependencyNotInstalled( "Pygame is not installed, run `pip install gym[classic_control]`" ) try: import matplotlib matplotlib.use("TkAgg") import matplotlib.pyplot as plt except ImportError: logger.warn("Matplotlib is not installed, run `pip install gym[other]`") matplotlib, plt = None, None class MissingKeysToAction(Exception): """Raised when the environment does not have a default ``keys_to_action`` mapping.""" class PlayableGame: """Wraps an environment allowing keyboard inputs to interact with the environment.""" def __init__( self, env: Env, keys_to_action: Optional[Dict[Tuple[int, ...], int]] = None, zoom: Optional[float] = None, ): """Wraps an environment with a dictionary of keyboard buttons to action and if to zoom in on the environment. Args: env: The environment to play keys_to_action: The dictionary of keyboard tuples and action value zoom: If to zoom in on the environment render """ if env.render_mode not in {"rgb_array", "rgb_array_list"}: logger.error( "PlayableGame wrapper works only with rgb_array and rgb_array_list render modes, " f"but your environment render_mode = {env.render_mode}." ) self.env = env self.relevant_keys = self._get_relevant_keys(keys_to_action) self.video_size = self._get_video_size(zoom) self.screen = pygame.display.set_mode(self.video_size) self.pressed_keys = [] self.running = True def _get_relevant_keys( self, keys_to_action: Optional[Dict[Tuple[int], int]] = None ) -> set: if keys_to_action is None: if hasattr(self.env, "get_keys_to_action"): keys_to_action = self.env.get_keys_to_action() elif hasattr(self.env.unwrapped, "get_keys_to_action"): keys_to_action = self.env.unwrapped.get_keys_to_action() else: raise MissingKeysToAction( f"{self.env.spec.id} does not have explicit key to action mapping, " "please specify one manually" ) assert isinstance(keys_to_action, dict) relevant_keys = set(sum((list(k) for k in keys_to_action.keys()), [])) return relevant_keys def _get_video_size(self, zoom: Optional[float] = None) -> Tuple[int, int]: rendered = self.env.render() if isinstance(rendered, List): rendered = rendered[-1] assert rendered is not None and isinstance(rendered, np.ndarray) video_size = (rendered.shape[1], rendered.shape[0]) if zoom is not None: video_size = (int(video_size[0] * zoom), int(video_size[1] * zoom)) return video_size def process_event(self, event: Event): """Processes a PyGame event. In particular, this function is used to keep track of which buttons are currently pressed and to exit the :func:`play` function when the PyGame window is closed. Args: event: The event to process """ if event.type == pygame.KEYDOWN: if event.key in self.relevant_keys: self.pressed_keys.append(event.key) elif event.key == pygame.K_ESCAPE: self.running = False elif event.type == pygame.KEYUP: if event.key in self.relevant_keys: self.pressed_keys.remove(event.key) elif event.type == pygame.QUIT: self.running = False elif event.type == VIDEORESIZE: self.video_size = event.size self.screen = pygame.display.set_mode(self.video_size) def display_arr( screen: Surface, arr: np.ndarray, video_size: Tuple[int, int], transpose: bool ): """Displays a numpy array on screen. Args: screen: The screen to show the array on arr: The array to show video_size: The video size of the screen transpose: If to transpose the array on the screen """ arr_min, arr_max = np.min(arr), np.max(arr) arr = 255.0 * (arr - arr_min) / (arr_max - arr_min) pyg_img = pygame.surfarray.make_surface(arr.swapaxes(0, 1) if transpose else arr) pyg_img = pygame.transform.scale(pyg_img, video_size) screen.blit(pyg_img, (0, 0)) def play( env: Env, transpose: Optional[bool] = True, fps: Optional[int] = None, zoom: Optional[float] = None, callback: Optional[Callable] = None, keys_to_action: Optional[Dict[Union[Tuple[Union[str, int]], str], ActType]] = None, seed: Optional[int] = None, noop: ActType = 0, ): """Allows one to play the game using keyboard. Example:: >>> import gym >>> from gym.utils.play import play >>> play(gym.make("CarRacing-v1", render_mode="rgb_array"), keys_to_action={ ... "w": np.array([0, 0.7, 0]), ... "a": np.array([-1, 0, 0]), ... "s": np.array([0, 0, 1]), ... "d": np.array([1, 0, 0]), ... "wa": np.array([-1, 0.7, 0]), ... "dw": np.array([1, 0.7, 0]), ... "ds": np.array([1, 0, 1]), ... "as": np.array([-1, 0, 1]), ... }, noop=np.array([0,0,0])) Above code works also if the environment is wrapped, so it's particularly useful in verifying that the frame-level preprocessing does not render the game unplayable. If you wish to plot real time statistics as you play, you can use :class:`gym.utils.play.PlayPlot`. Here's a sample code for plotting the reward for last 150 steps. >>> def callback(obs_t, obs_tp1, action, rew, terminated, truncated, info): ... return [rew,] >>> plotter = PlayPlot(callback, 150, ["reward"]) >>> play(gym.make("ALE/AirRaid-v5"), callback=plotter.callback) Args: env: Environment to use for playing. transpose: If this is ``True``, the output of observation is transposed. Defaults to ``True``. fps: Maximum number of steps of the environment executed every second. If ``None`` (the default), ``env.metadata["render_fps""]`` (or 30, if the environment does not specify "render_fps") is used. zoom: Zoom the observation in, ``zoom`` amount, should be positive float callback: If a callback is provided, it will be executed after every step. It takes the following input: obs_t: observation before performing action obs_tp1: observation after performing action action: action that was executed rew: reward that was received terminated: whether the environment is terminated or not truncated: whether the environment is truncated or not info: debug info keys_to_action: Mapping from keys pressed to action performed. Different formats are supported: Key combinations can either be expressed as a tuple of unicode code points of the keys, as a tuple of characters, or as a string where each character of the string represents one key. For example if pressing 'w' and space at the same time is supposed to trigger action number 2 then ``key_to_action`` dict could look like this: >>> { ... # ... ... (ord('w'), ord(' ')): 2 ... # ... ... } or like this: >>> { ... # ... ... ("w", " "): 2 ... # ... ... } or like this: >>> { ... # ... ... "w ": 2 ... # ... ... } If ``None``, default ``key_to_action`` mapping for that environment is used, if provided. seed: Random seed used when resetting the environment. If None, no seed is used. noop: The action used when no key input has been entered, or the entered key combination is unknown. """ env.reset(seed=seed) if keys_to_action is None: if hasattr(env, "get_keys_to_action"): keys_to_action = env.get_keys_to_action() elif hasattr(env.unwrapped, "get_keys_to_action"): keys_to_action = env.unwrapped.get_keys_to_action() else: raise MissingKeysToAction( f"{env.spec.id} does not have explicit key to action mapping, " "please specify one manually" ) assert keys_to_action is not None key_code_to_action = {} for key_combination, action in keys_to_action.items(): key_code = tuple( sorted(ord(key) if isinstance(key, str) else key for key in key_combination) ) key_code_to_action[key_code] = action game = PlayableGame(env, key_code_to_action, zoom) if fps is None: fps = env.metadata.get("render_fps", 30) done, obs = True, None clock = pygame.time.Clock() while game.running: if done: done = False obs = env.reset(seed=seed) else: action = key_code_to_action.get(tuple(sorted(game.pressed_keys)), noop) prev_obs = obs obs, rew, terminated, truncated, info = env.step(action) done = terminated or truncated if callback is not None: callback(prev_obs, obs, action, rew, terminated, truncated, info) if obs is not None: rendered = env.render() if isinstance(rendered, List): rendered = rendered[-1] assert rendered is not None and isinstance(rendered, np.ndarray) display_arr( game.screen, rendered, transpose=transpose, video_size=game.video_size ) # process pygame events for event in pygame.event.get(): game.process_event(event) pygame.display.flip() clock.tick(fps) pygame.quit() class PlayPlot: """Provides a callback to create live plots of arbitrary metrics when using :func:`play`. This class is instantiated with a function that accepts information about a single environment transition: - obs_t: observation before performing action - obs_tp1: observation after performing action - action: action that was executed - rew: reward that was received - terminated: whether the environment is terminated or not - truncated: whether the environment is truncated or not - info: debug info It should return a list of metrics that are computed from this data. For instance, the function may look like this:: >>> def compute_metrics(obs_t, obs_tp, action, reward, terminated, truncated, info): ... return [reward, info["cumulative_reward"], np.linalg.norm(action)] :class:`PlayPlot` provides the method :meth:`callback` which will pass its arguments along to that function and uses the returned values to update live plots of the metrics. Typically, this :meth:`callback` will be used in conjunction with :func:`play` to see how the metrics evolve as you play:: >>> plotter = PlayPlot(compute_metrics, horizon_timesteps=200, ... plot_names=["Immediate Rew.", "Cumulative Rew.", "Action Magnitude"]) >>> play(your_env, callback=plotter.callback) """ def __init__( self, callback: callable, horizon_timesteps: int, plot_names: List[str] ): """Constructor of :class:`PlayPlot`. The function ``callback`` that is passed to this constructor should return a list of metrics that is of length ``len(plot_names)``. Args: callback: Function that computes metrics from environment transitions horizon_timesteps: The time horizon used for the live plots plot_names: List of plot titles Raises: DependencyNotInstalled: If matplotlib is not installed """ deprecation( "`PlayPlot` is marked as deprecated and will be removed in the near future." ) self.data_callback = callback self.horizon_timesteps = horizon_timesteps self.plot_names = plot_names if plt is None: raise DependencyNotInstalled( "matplotlib is not installed, run `pip install gym[other]`" ) num_plots = len(self.plot_names) self.fig, self.ax = plt.subplots(num_plots) if num_plots == 1: self.ax = [self.ax] for axis, name in zip(self.ax, plot_names): axis.set_title(name) self.t = 0 self.cur_plot: List[Optional[plt.Axes]] = [None for _ in range(num_plots)] self.data = [deque(maxlen=horizon_timesteps) for _ in range(num_plots)] def callback( self, obs_t: ObsType, obs_tp1: ObsType, action: ActType, rew: float, terminated: bool, truncated: bool, info: dict, ): """The callback that calls the provided data callback and adds the data to the plots. Args: obs_t: The observation at time step t obs_tp1: The observation at time step t+1 action: The action rew: The reward terminated: If the environment is terminated truncated: If the environment is truncated info: The information from the environment """ points = self.data_callback( obs_t, obs_tp1, action, rew, terminated, truncated, info ) for point, data_series in zip(points, self.data): data_series.append(point) self.t += 1 xmin, xmax = max(0, self.t - self.horizon_timesteps), self.t for i, plot in enumerate(self.cur_plot): if plot is not None: plot.remove() self.cur_plot[i] = self.ax[i].scatter( range(xmin, xmax), list(self.data[i]), c="blue" ) self.ax[i].set_xlim(xmin, xmax) if plt is None: raise DependencyNotInstalled( "matplotlib is not installed, run `pip install gym[other]`" ) plt.pause(0.000001)