-
-
Notifications
You must be signed in to change notification settings - Fork 712
/
play.py
417 lines (347 loc) · 15.9 KB
/
play.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
"""Utilities of visualising an environment."""
from __future__ import annotations
from collections import deque
from typing import Callable, List
import numpy as np
import gymnasium as gym
from gymnasium import Env, logger
from gymnasium.core import ActType, ObsType
from gymnasium.error import DependencyNotInstalled
try:
import pygame
from pygame import Surface
from pygame.event import Event
except ImportError as e:
raise gym.error.DependencyNotInstalled(
'pygame is not installed, run `pip install "gymnasium[classic_control]"`'
) from e
try:
import matplotlib
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt
except ImportError:
logger.warn('matplotlib is not installed, run `pip install "gymnasium[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: dict[tuple[int, ...], int] | None = None,
zoom: float | None = 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"}:
raise ValueError(
"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 is the size of the video that is being displayed.
# The window size may be larger, in that case we will add black bars
self.video_size = self._get_video_size(zoom)
self.screen = pygame.display.set_mode(self.video_size, pygame.RESIZABLE)
self.pressed_keys = []
self.running = True
def _get_relevant_keys(
self, keys_to_action: dict[tuple[int], int] | None = None
) -> set:
if keys_to_action is None:
if self.env.has_wrapper_attr("get_keys_to_action"):
keys_to_action = self.env.get_wrapper_attr("get_keys_to_action")()
else:
assert self.env.spec is not None
raise MissingKeysToAction(
f"{self.env.spec.id} does not have explicit key to action mapping, "
"please specify one manually, `play(env, keys_to_action=...)`"
)
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: float | None = 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 == pygame.WINDOWRESIZED:
# Compute the maximum video size that fits into the new window
scale_width = event.x / self.video_size[0]
scale_height = event.y / self.video_size[1]
scale = min(scale_height, scale_width)
self.video_size = (scale * self.video_size[0], scale * self.video_size[1])
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
"""
assert isinstance(arr, np.ndarray) and arr.dtype == np.uint8
pyg_img = pygame.surfarray.make_surface(arr.swapaxes(0, 1) if transpose else arr)
pyg_img = pygame.transform.scale(pyg_img, video_size)
# We might have to add black bars if surface_size is larger than video_size
surface_size = screen.get_size()
width_offset = (surface_size[0] - video_size[0]) / 2
height_offset = (surface_size[1] - video_size[1]) / 2
screen.fill((0, 0, 0))
screen.blit(pyg_img, (width_offset, height_offset))
def play(
env: Env,
transpose: bool | None = True,
fps: int | None = None,
zoom: float | None = None,
callback: Callable | None = None,
keys_to_action: dict[tuple[str | int, ...] | str | int, ActType] | None = None,
seed: int | None = None,
noop: ActType = 0,
wait_on_player: bool = False,
):
"""Allows the user to play the environment using a keyboard.
If playing in a turn-based environment, set wait_on_player to True.
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:
>>> key_to_action = {
... # ...
... (ord('w'), ord(' ')): 2
... # ...
... }
or like this:
>>> key_to_action = {
... # ...
... ("w", " "): 2
... # ...
... }
or like this:
>>> key_to_action = {
... # ...
... "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.
wait_on_player: Play should wait for a user action
Example:
>>> import gymnasium as gym
>>> import numpy as np
>>> from gymnasium.utils.play import play
>>> play(gym.make("CarRacing-v2", render_mode="rgb_array"), # doctest: +SKIP
... 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:`PlayPlot`. Here's a sample code for plotting the reward
for last 150 steps.
>>> from gymnasium.utils.play import PlayPlot, play
>>> def callback(obs_t, obs_tp1, action, rew, terminated, truncated, info):
... return [rew,]
>>> plotter = PlayPlot(callback, 150, ["reward"]) # doctest: +SKIP
>>> play(gym.make("CartPole-v1"), callback=plotter.callback) # doctest: +SKIP
"""
env.reset(seed=seed)
if keys_to_action is None:
if env.has_wrapper_attr("get_keys_to_action"):
keys_to_action = env.get_wrapper_attr("get_keys_to_action")()
else:
assert env.spec is not None
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
# validate the `keys_to_action` set provided
assert isinstance(keys_to_action, dict)
for key, action in keys_to_action.items():
if isinstance(key, tuple):
assert len(key) > 0
assert all(isinstance(k, (str, int)) for k in key)
else:
assert isinstance(key, (str, int))
assert action in env.action_space
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)
elif wait_on_player is False or len(game.pressed_keys) > 0:
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, # doctest: +SKIP
... plot_names=["Immediate Rew.", "Cumulative Rew.", "Action Magnitude"])
>>> play(your_env, callback=plotter.callback) # doctest: +SKIP
"""
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
"""
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 "gymnasium[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[plt.Axes | None] = [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 "gymnasium[other]"`'
)
plt.pause(0.000001)