Cope 2.5.0
My personal "standard library" of all the generally useful code I've written for various projects over the years
Loading...
Searching...
No Matches
gym.py
1"""
2Functions & classes that extend the gymnasium library
3"""
4__version__ = '1.0.0'
5
6from .misc import RedirectStd
7from sys import exit
8from abc import ABC
9from time import sleep, time as now
10from .imports import lazy_import, ensure_imported
11# try:
12
13# import gymnasium as gym
14gym = lazy_import('gymnasium')
15# Don't print the annoying "welcome from the pygame community!" message
16# with RedirectStd():
17pygame = lazy_import('pygame')
18# import pygame
19# from pygame import Rect
20# except:
21# # So it still gives an error, but only if we try to use it
22# class SimpleGym:
23# def __init__(self):
24# raise ImportError("gymnasium or pygame not installed. Please install both before importing SimpleGym.")
25# else:
26if gym:
27 class SimpleGym(gym.Env, ABC):
28 """ A simplified Gymnasium enviorment that uses pygame and handles some stuff for you, like rendering
29 keeping track of steps, returning the right things from the right functions, event handling,
30 including some default keyboard shortcuts, and some debugging features. Includes easy ways
31 to print to the screen.
32
33 **Rendering**
34 By default, it's set to render the enviorment according to the function render_pygame, which
35 should render everything to self.surf. If you would like to specify other rendering methods,
36 define them as `render_{method_name}`, and render() will handle it for you. There's no need
37 to manually overwrite the render() method.
38
39 **Printing**
40 There are 3 ways to print to the screen:
41 `show_vars`: accessable either via the constructor or as a member
42 This is a dictionary of {name: member} of members you want to have printed
43 on the screen. The keys can be any string, and the values must be valid members
44 of this class. The screen is updated as the member changes.
45 `show_strings`: accessable either via the constructor, as a member, or the show() function
46 This is a list of strings that are printed to the screen. They won't change.
47 Useful for showing config options and the like. They are always printed first.
48 `print`: a dictionary member. The keys are arbitrary and not printed. The values
49 are printed to the screen. The reason it's a dictionary and not a list is simply
50 for easy indexing: you can change this in the middle of the running loop, and
51 it will get added to the screen. Just make sure to reuse the keys.
52 Attempting to set an item on an instance will also set it to print
53 (i.e. `env['a'] = 'string'` is the same as `env.print['a'] = 'string'`)
54
55 Default key handling:
56 q closes the window
57 space toggles pause
58 i increments a single frame
59 u runs the debug_button()
60 r runs reset() immediately
61 f toggles whether we're limiting ourselves to FPS or not
62 h shows a help menu on screen
63 >/< increase and decrease the FPS
64
65 In order to use this effectively, you need to overload:
66 __init__(), if you want to add any members
67 _get_obs()
68 _get_reward()
69 _step(action), don't overload step(), step() calls _step
70 _reset(seed=None, options=None), if you need any custom reset code (don't overload reset())
71 render_pygame(), render to self.surf
72 and optionally:
73 _get_terminated(), if you want to use custom terminated criteria other than just max_steps
74 You *don't* need to call super()._get_terminated()
75 _get_info(), if you want to include info
76 handle_event(event), for handling events
77 debug_button(), for debugging when you press the u key
78 _get_truncated(), if you want to include truncated
79
80 Helpful members provided:
81 `width`, `height`: the dimentions of the screen
82 `size`: set to self.width, for compatibility's sake
83 `just_reset`: is set to True by reset() and False by step()
84 `steps`: the number of steps in the current episode
85 `total_steps`: the total number of steps taken
86 `reset_count`: a count of the number of times reset() has been called
87 `surf`: the surface to draw to in render_pygame()
88 `paused`: True if paused
89 `increment`: set to True internally to denote a single step while paused. step() sets
90 to False
91 """
92
93 FPS_STEP = 2
94 SHOW_HELP_FADE_TIME = 10
95 FONT_SIZE = 10
96 HELP_TEXT = """
97 q: close window
98 space: toggle pause
99 i: increment a single frame
100 u: run debug_button()
101 r: reset
102 f: toggle limit FPS
103 </>: change FPS
104 h: show/hide this menu
105 """
106
107 def __init__(self,
108 max_steps=None,
109 screen_size=300,
110 fps=None,
111 name='SimpleGym Enviorment',
112 show_vars={},
113 show_strings=[],
114 start_paused=False,
115 render_mode='pygame',
116 assert_valid_action=True,
117 background_color=(20, 20, 20),
118 print_color=(200, 20, 20, 0),
119 show_events=False,
120 verbose=True,
121 ):
122 """ This should be called first, if you want to use the members like self.size
123 Parameters:
124 `max_steps`: if positive, sets the maximum number of steps before the env resets
125 itself. If None or negative, no limit
126 `screen_size`: the size of the pygame window. Can be a 2 item tuple of (width, height)
127 or a single int if the window is to be square
128 `fps`: controls how fast the simulation runs. Set to negative or None to have no limit
129 `name`: the name of the enviorment shown on the window title
130 `show_vars`: a dictionary of {name: member} of members you want to have printed
131 on the screen. The keys can be any string, and the values must be valid members
132 of this class
133 `show_strings`: a list of strings you want to have printed on the screen
134 `start_paused`: self-explanitory
135 `show_events`: prints events, other than mouse movements, for debugging purpouses
136 `render_mode`: part of the gymnasium specification. Must be either None or 'pygame',
137 unless you manually override the render() method
138 `assert_valid_action`: ensures that actions given to step() are within the action_space
139 `background_color`: a 3 item tuple specifying the background color
140 `print_color`: a 4 item tuple (the 4th index being alpha) specifying the color of
141 the extra data printed to the screen
142 `verbose`: when set to True, it simply adds `fps`, `reset_count`, `steps`, `total_steps`
143 to `show_vars`. Also shows the help menu for the first few seconds
144 """
145 self.metadata = {"render_modes": list({'pygame', render_mode}), "render_fps": fps}
146 assert render_mode is None or render_mode in self.metadata["render_modes"], render_mode
147
148 self.background_color = background_color
149 self.print_color = print_color
150 self.name = name
151 self.show_events = show_events
152 self.fps = self.FPS = fps
153
154 self.max_steps = max_steps
155 self.steps = 0
156 self.total_steps = 0
157 self.paused = start_paused
158 self.increment = False
159 self.just_reset = False
160 self.reset_count = 0
161
162 self.render_mode = render_mode
163 self.screen_size = screen_size if isinstance(screen_size, (tuple, list)) else (screen_size, screen_size)
164 self.width, self.height = self.screen_size
165 self.size = self.width
166 self.screen = None
167 self.surf = None
168 self.font = None
169
170 self.print = {}
171 self.show_strings = show_strings
172 self.show_vars = show_vars
173 if verbose:
174 self.show_vars['FPS'] = 'fps'
175 self.show_vars['Episode'] = 'reset_count'
176 self.show_vars['Step'] = 'steps'
177 self.show_vars['Total Steps'] = 'total_steps'
178
179 # We want to ensure that reset gets called before step. Nowhere else does this get set to False
180 self._previous_step_called = None
181 self._assert_valid_action = assert_valid_action
182 self._prints_surf = None
183 self._original_fps = fps
184 self._show_help = verbose
185 self._rendered_helps = None
186
187 def _get_obs(self):
188 raise NotImplementedError
189
190 def _get_info(self):
191 return {}
192
193 def _get_truncated(self):
194 return False
195
196 def __default_get_terminated(self):
197 """ Called internally so max_steps still works even if _get_terminated is overloaded """
198 term = False
199 if self.max_steps is not None and self.max_steps > 0 and self.steps > self.max_steps:
200 print(term)
201 term = True
202
203 return self._get_terminated() or term
204
206 """ By default this just terminates after max_steps have been reached """
207 return False
208
209 def _get_reward(self):
210 raise NotImplementedError
211
212 def reset(self, seed=None, options=None):
213 """ This sets the self.np_random to use the seed given, resets the steps, and returns the
214 observation and info. Needs to be called first, if you're depending on self.np_random,
215 or steps equalling 0, but also needs to return what this returns.
216 """
217 # We need the following line to seed self.np_random
218 super().reset(seed=seed)
219
220 self._reset(seed=seed, options=options)
221
222 self.steps = 0
223 self.just_reset = True
224 self.reset_count += 1
225
226 return self._get_obs(), self._get_info()
227
228 def _reset(seed=None, options=None):
229 raise NotImplementedError()
230
231 def step(self, action):
232 """ Call this last, and return it """
233 assert self.reset_count > 0, "step() called before reset(). reset() must be called first."
234
235 # If it's paused, don't bother checking if it's a valid action
236 if self._assert_valid_action and not self.paused or self.increment:
237 assert self.action_space.contains(action), "Action given not within action_space"
238
239 if self.paused and not self.increment:
240 return self._get_obs(), self._get_reward(), self.__default_get_terminated(), self._get_truncated(), self._get_info()
241
242 self._step(action)
243
244 if not self.paused or self.increment:
245 self.steps += 1
246 self.total_steps += 1
247 self.just_reset = False
248 self.increment = False
249
250 if self.fps is not None and self.fps > 0 and self._previous_step_called is not None:
251 # Sleep for the amount of time until we're supposed to call step next
252 wait_for = (1/self.fps) - (now() - self._previous_step_called)
253 # If calculations elsewhere took so long that we've already past the next frame time,
254 # don't sleep, just run
255 if wait_for > 0:
256 sleep(wait_for)
257
258 self._previous_step_called = now()
259 return self._get_obs(), self._get_reward(), self.__default_get_terminated(), self._get_truncated(), self._get_info()
260
261 def _step(action):
262 raise NotImplementedError()
263
264 @ensure_imported(pygame)
265 def _init_pygame(self):
266 if self.screen is None:
267 pygame.init()
268 pygame.display.init()
269 pygame.display.set_caption(self.name)
270 self.screen = pygame.display.set_mode(self.screen_size)
271
272 if self.font is None:
273 self.font = pygame.font.SysFont("Verdana", self.FONT_SIZE)
274
275 if self.surf is None:
276 self.surf = pygame.Surface(self.screen_size)
277 self.surf.convert()
278 # self.surf.fill((255, 255, 255))
279
280 if self._prints_surf is None:
281 self._prints_surf = pygame.Surface(self.screen_size)
282 self._prints_surf.convert()
283 self._prints_surf.fill(self.background_color)
284
285 if self._rendered_helps is None:
286 self._rendered_helps = [
287 self.font.render(line, True, self.print_color)
288 for line in self.HELP_TEXT.splitlines()
289 ]
290
291 def render(self):
292 if self.render_mode == 'pygame':
293 self._init_pygame()
294 # Fill the background
295 self.surf.fill(self.background_color)
296
297 self.render_pygame()
298
299 # The strings in show_strings should come first
300 strings = self.show_strings.copy()
301 # Get the texts from self.show_vars
302 length = len(max(self.show_vars.keys(), key=len))
303 strings += [
304 f'{name}: {" "*(length - len(name))} {getattr(self, var, f"{var} is not a member")}'
305 for name, var in self.show_vars.items()
306 ]
307
308 # Add the texts from self.prints
309 strings += list(self.print.values())
310
311 # Draw all the text onto the surface
312 for offset, string in enumerate(strings):
313 self.surf.blit(self.font.render(string, True, self.print_color), (5, 5 + offset*(self.FONT_SIZE + 2)))
314
315 if self._show_help:
316 for offset, string in enumerate(self._rendered_helps):
317 max_width = max(self._rendered_helps, key=lambda h: h.get_size()[0]).get_size()[0]
318 self.surf.blit(string, (self.width - max_width, offset*(self.FONT_SIZE + 2)))
319
320 # I don't remember what this does, but I think it's important
321 self.surf = pygame.transform.scale(self.surf, self.screen_size)
322
323 # Display to screen
324 self.screen.blit(self.surf, (0, 0))
325 self._handle_events()
326 pygame.display.flip()
327
328 else:
329 if hasattr(self, f'render_{self.render_mode}'):
330 getattr(self, f'render_{self.render_mode}')()
331 else:
332 raise AttributeError(f"No render_{self.render_mode} function provided")
333
334 def debug_button(self):
335 pass
336
337 def _handle_events(self):
338 for e in pygame.event.get():
339 match e.type:
340 case pygame.QUIT:
341 self.close()
342 exit(0)
343 case pygame.KEYDOWN:
344 if e.key == pygame.K_ESCAPE:
345 self.close()
346 exit(0)
347 elif e.key == pygame.K_SPACE:
348 self.paused = not self.paused
349
350 match e.unicode:
351 case 'q':
352 self.close()
353 exit(0)
354 case 'u':
355 self.debug_button()
356 case 'i':
357 self.increment = True
358 # If it's not paused, make it paused
359 self.paused = True
360 case 'r':
361 self.reset()
362 case 'f':
363 if self.fps is None:
364 self.fps = self._original_fps
365 else:
366 self._original_fps = self.fps
367 self.fps = None
368 case '>' | '.':
369 self.fps += self.FPS_STEP
370 case '<' | ',':
371 self.fps -= self.FPS_STEP
372 case 'h':
373 self._show_help = not self._show_help
374 case _:
375 self.handle_event(e)
376 case _:
377 self.handle_event(e)
378 if self.show_events and e.type != pygame.MOUSEMOTION:
379 print(e)
380 pygame.event.pump()
381
382 def handle_event(self, event):
383 pass
384
385 def show(self, string):
386 """ Sets a given string to be shown on the pygame window """
387 self.show_strings.append(string)
388
389 def __setitem__(self, index, string):
390 self.print[index] = string
391
392 def close(self):
393 if self.screen is not None:
394 pygame.display.quit()
395 pygame.quit()
396 self.screen = None
397 self.font = None
398
399else:
400 # So it still gives an error, but only if we try to use it
401 class SimpleGym:
402 def __new__(*args, **kwargs):
403 raise ImportError("gymnasium not installed. Run `pip install gymnasium` before importing SimpleGym.")
def _get_info(self)
Definition: gym.py:190
def _init_pygame(self)
Definition: gym.py:265
def _step(action)
Definition: gym.py:261
def _reset(seed=None, options=None)
Definition: gym.py:228
def reset(self, seed=None, options=None)
This sets the self.np_random to use the seed given, resets the steps, and returns the observation and...
Definition: gym.py:212
def show(self, string)
Sets a given string to be shown on the pygame window.
Definition: gym.py:385
def _get_truncated(self)
Definition: gym.py:193
def step(self, action)
Call this last, and return it.
Definition: gym.py:231
def _get_terminated(self)
By default this just terminates after max_steps have been reached.
Definition: gym.py:205
def _handle_events(self)
Definition: gym.py:337
str HELP_TEXT
Definition: gym.py:96
def _get_obs(self)
Definition: gym.py:187
def _get_reward(self)
Definition: gym.py:209
def __init__(self, max_steps=None, screen_size=300, fps=None, name='SimpleGym Enviorment', show_vars={}, show_strings=[], start_paused=False, render_mode='pygame', assert_valid_action=True, background_color=(20, 20, 20), print_color=(200, 20, 20, 0), show_events=False, verbose=True)
This should be called first, if you want to use the members like self.size Parameters: max_steps: if ...
Definition: gym.py:121
int FONT_SIZE
Definition: gym.py:95
A simplified Gymnasium enviorment that uses pygame and handles some stuff for you,...
Definition: gym.py:27
def __default_get_terminated(self)
Definition: gym.py:196