smb2-gym 0.2.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (30) hide show
  1. smb2_gym-0.2.0/MANIFEST.in +4 -0
  2. smb2_gym-0.2.0/PKG-INFO +209 -0
  3. smb2_gym-0.2.0/README.md +178 -0
  4. smb2_gym-0.2.0/pyproject.toml +86 -0
  5. smb2_gym-0.2.0/setup.cfg +4 -0
  6. smb2_gym-0.2.0/smb2_gym/__init__.py +35 -0
  7. smb2_gym-0.2.0/smb2_gym/_nes/prg0/saves/luigi/1-1.sav +0 -0
  8. smb2_gym-0.2.0/smb2_gym/_nes/prg0/saves/mario/1-1.sav +0 -0
  9. smb2_gym-0.2.0/smb2_gym/_nes/prg0/saves/peach/1-1.sav +0 -0
  10. smb2_gym-0.2.0/smb2_gym/_nes/prg0/saves/toad/1-1.sav +0 -0
  11. smb2_gym-0.2.0/smb2_gym/_nes/prg0/super_mario_bros_2_prg0.nes +0 -0
  12. smb2_gym-0.2.0/smb2_gym/_nes/prg0_edited/saves/easy_combined_curriculum.sav +0 -0
  13. smb2_gym-0.2.0/smb2_gym/_nes/prg0_edited/super_mario_bros_2_prg0_edited.nes +0 -0
  14. smb2_gym-0.2.0/smb2_gym/actions.py +168 -0
  15. smb2_gym-0.2.0/smb2_gym/app/__init__.py +138 -0
  16. smb2_gym-0.2.0/smb2_gym/app/info_display.py +123 -0
  17. smb2_gym-0.2.0/smb2_gym/app/keyboard.py +25 -0
  18. smb2_gym-0.2.0/smb2_gym/app/rendering.py +21 -0
  19. smb2_gym-0.2.0/smb2_gym/constants.py +155 -0
  20. smb2_gym-0.2.0/smb2_gym/play_human.py +313 -0
  21. smb2_gym-0.2.0/smb2_gym/py.typed +0 -0
  22. smb2_gym-0.2.0/smb2_gym/smb2_env.py +786 -0
  23. smb2_gym-0.2.0/smb2_gym.egg-info/PKG-INFO +209 -0
  24. smb2_gym-0.2.0/smb2_gym.egg-info/SOURCES.txt +28 -0
  25. smb2_gym-0.2.0/smb2_gym.egg-info/dependency_links.txt +1 -0
  26. smb2_gym-0.2.0/smb2_gym.egg-info/entry_points.txt +2 -0
  27. smb2_gym-0.2.0/smb2_gym.egg-info/requires.txt +10 -0
  28. smb2_gym-0.2.0/smb2_gym.egg-info/top_level.txt +1 -0
  29. smb2_gym-0.2.0/tests/test_cli_options.py +54 -0
  30. smb2_gym-0.2.0/tests/test_random_actions.py +76 -0
@@ -0,0 +1,4 @@
1
+ include README.md
2
+ include LICENSE
3
+ recursive-include smb2_gym *.py
4
+ recursive-include smb2_gym/_nes *.whl *.nes *.sav
@@ -0,0 +1,209 @@
1
+ Metadata-Version: 2.4
2
+ Name: smb2-gym
3
+ Version: 0.2.0
4
+ Summary: Super Mario Bros 2 (Europe) Gymnasium Environment
5
+ Author: Alexei Quick
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/alexeiquickcode/smb2-gym
8
+ Project-URL: Bug Tracker, https://github.com/alexeiquickcode/smb2-gym/issues
9
+ Classifier: Development Status :: 3 - Alpha
10
+ Classifier: Intended Audience :: Science/Research
11
+ Classifier: Operating System :: OS Independent
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Programming Language :: Python :: 3.8
14
+ Classifier: Programming Language :: Python :: 3.9
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
20
+ Requires-Python: >=3.8
21
+ Description-Content-Type: text/markdown
22
+ Requires-Dist: gymnasium>=0.29.0
23
+ Requires-Dist: numpy>=1.21.0
24
+ Requires-Dist: pygame>=2.1.0
25
+ Requires-Dist: tetanes-py>=0.1.1
26
+ Provides-Extra: dev
27
+ Requires-Dist: pytest; extra == "dev"
28
+ Requires-Dist: isort; extra == "dev"
29
+ Requires-Dist: yapf; extra == "dev"
30
+ Requires-Dist: basedpyright; extra == "dev"
31
+
32
+ # smb2-gym
33
+
34
+ [![Python](https://img.shields.io/pypi/pyversions/smb2-gym)](https://pypi.org/project/smb2-gym/)
35
+ [![PyPI](https://img.shields.io/pypi/v/smb2-gym)](https://pypi.org/project/smb2-gym/)
36
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
37
+
38
+ A Gymnasium environment for Super Mario Bros 2 (Europe/Doki Doki Panic version) using TetaNES emulator Python bindings. Perfect for reinforcement learning experiments and research.
39
+
40
+ **Features:**
41
+ - Curated action sets for faster training (`simple`, `complex`)
42
+ - Comprehensive game state via info dict (40+ properties)
43
+ - Multiple initialization modes (character/level, custom ROMs, save states)
44
+ - Human-playable interface with keyboard controls
45
+
46
+ A full list of the available RAM map properties is available at [Data Crystal](https://datacrystal.tcrf.net/wiki/Super_Mario_Bros._2_(NES)/RAM_map).
47
+
48
+ ## Installation
49
+
50
+ ```bash
51
+ pip install smb2-gym
52
+ ```
53
+
54
+ ## Quick Start
55
+
56
+ ### Basic Usage
57
+
58
+ ```python
59
+ from smb2_gym import SuperMarioBros2Env
60
+ from smb2_gym.app import InitConfig
61
+
62
+ # Create environment with character/level mode
63
+ config = InitConfig(level="1-1", character="luigi")
64
+ env = SuperMarioBros2Env(
65
+ init_config=config,
66
+ render_mode="human", # "human" or None
67
+ action_type="simple" # "simple" (11), "complex" (16), or "all" (256)
68
+ )
69
+
70
+ # Reset environment
71
+ obs, info = env.reset()
72
+
73
+ # Run game loop
74
+ for _ in range(1000):
75
+ action = env.action_space.sample()
76
+ obs, reward, terminated, truncated, info = env.step(action)
77
+
78
+ # Access game state from info dict
79
+ print(f"Lives: {info['life']}, Hearts: {info['hearts']}")
80
+ print(f"Position: ({info['x_pos_global']}, {info['y_pos_global']})")
81
+
82
+ if terminated or truncated:
83
+ obs, info = env.reset()
84
+
85
+ env.close()
86
+ ```
87
+
88
+ ### Initialization Modes
89
+
90
+ ```python
91
+ from smb2_gym.app import InitConfig
92
+
93
+ # 1. Character/Level mode (default)
94
+ config = InitConfig(level="1-1", character="peach")
95
+
96
+ # 2. Built-in ROM variant mode
97
+ config = InitConfig(rom="prg0", save_state="level_1_1.sav")
98
+
99
+ # 3. Custom ROM mode
100
+ config = InitConfig(
101
+ rom_path="/path/to/your/smb2.nes",
102
+ save_state_path="/path/to/save.sav" # Optional
103
+ )
104
+ ```
105
+
106
+ ### Custom Reward Function
107
+
108
+ ```python
109
+ from smb2_gym import SuperMarioBros2Env
110
+ from smb2_gym.app import InitConfig
111
+
112
+ class CustomSMB2Env(SuperMarioBros2Env):
113
+ def step(self, action):
114
+ obs, reward, terminated, truncated, info = super().step(action)
115
+
116
+ # Custom reward based on x-position progress
117
+ reward = info['x_pos_global'] / 100.0
118
+
119
+ # Bonus for collecting cherries
120
+ reward += info['cherries'] * 10
121
+
122
+ # Bonus for hearts
123
+ reward += info['hearts'] * 5
124
+
125
+ # Penalty for losing a life
126
+ if info.get('life_lost'):
127
+ reward -= 100
128
+
129
+ return obs, reward, terminated, truncated, info
130
+
131
+ # Usage
132
+ config = InitConfig(level="1-1", character="luigi")
133
+ env = CustomSMB2Env(init_config=config, action_type="simple")
134
+ ```
135
+
136
+ ## Play as Human
137
+
138
+ The package includes a human-playable interface with multiple initialization modes:
139
+
140
+ ### Character/Level Mode (Default)
141
+ ```bash
142
+ # Play as Luigi on level 1-1
143
+ smb2-play --level 1-1 --char luigi --scale 3
144
+
145
+ # TODO: Yet to implement all levels states
146
+ # Play as Peach on level 2-3
147
+ smb2-play --level 2-3 --char peach
148
+ ```
149
+
150
+ ### Built-in ROM Variant Mode
151
+ ```bash
152
+ # Use specific ROM variant with save state
153
+ smb2-play --rom prg0_edited --save-state easy_combined_curriculum.sav
154
+ ```
155
+
156
+ ### Custom ROM Mode
157
+ ```bash
158
+ # Use your own ROM file
159
+ smb2-play --custom-rom /path/to/smb2.nes
160
+
161
+ # Use custom ROM with save state
162
+ smb2-play --custom-rom /path/to/smb2.nes --custom-state /path/to/save.sav
163
+
164
+ # Start from beginning without save state
165
+ smb2-play --custom-rom /path/to/smb2.nes --no-save-state
166
+ ```
167
+
168
+ ### Controls
169
+
170
+ **Primary Controls:**
171
+ - Arrow Keys: Move
172
+ - Z: A button (Jump)
173
+ - X: B button (Pick up/Throw)
174
+ - Enter: Start
175
+ - Right Shift: Select
176
+ - P: Pause
177
+ - R: Reset
178
+ - ESC: Quit
179
+
180
+ **Alternative Controls:**
181
+ - WASD: Move
182
+ - J: A button
183
+ - K: B button
184
+
185
+ **Save States:**
186
+ - F5: Save state
187
+ - F9: Load state
188
+
189
+ ### CLI Options
190
+
191
+ **Character/Level Mode:**
192
+ - `--level`: Level to play (1-1 through 7-2, default: 1-1)
193
+ - `--char`: Character (mario, luigi, peach, toad, default: luigi)
194
+
195
+ **Built-in ROM Mode:**
196
+ - `--rom`: ROM variant (prg0, prg0_edited)
197
+ - `--save-state`: Save state filename
198
+
199
+ **Custom ROM Mode:**
200
+ - `--custom-rom`: Path to custom ROM file
201
+ - `--custom-state`: Path to custom save state (optional)
202
+ - `--no-save-state`: Start from beginning without loading save state
203
+
204
+ **Display:**
205
+ - `--scale`: Display scale factor (1-4, default: 3)
206
+
207
+ ## Disclaimer
208
+
209
+ This project is for educational and research purposes only. Users must provide their own legally obtained ROM files.
@@ -0,0 +1,178 @@
1
+ # smb2-gym
2
+
3
+ [![Python](https://img.shields.io/pypi/pyversions/smb2-gym)](https://pypi.org/project/smb2-gym/)
4
+ [![PyPI](https://img.shields.io/pypi/v/smb2-gym)](https://pypi.org/project/smb2-gym/)
5
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
6
+
7
+ A Gymnasium environment for Super Mario Bros 2 (Europe/Doki Doki Panic version) using TetaNES emulator Python bindings. Perfect for reinforcement learning experiments and research.
8
+
9
+ **Features:**
10
+ - Curated action sets for faster training (`simple`, `complex`)
11
+ - Comprehensive game state via info dict (40+ properties)
12
+ - Multiple initialization modes (character/level, custom ROMs, save states)
13
+ - Human-playable interface with keyboard controls
14
+
15
+ A full list of the available RAM map properties is available at [Data Crystal](https://datacrystal.tcrf.net/wiki/Super_Mario_Bros._2_(NES)/RAM_map).
16
+
17
+ ## Installation
18
+
19
+ ```bash
20
+ pip install smb2-gym
21
+ ```
22
+
23
+ ## Quick Start
24
+
25
+ ### Basic Usage
26
+
27
+ ```python
28
+ from smb2_gym import SuperMarioBros2Env
29
+ from smb2_gym.app import InitConfig
30
+
31
+ # Create environment with character/level mode
32
+ config = InitConfig(level="1-1", character="luigi")
33
+ env = SuperMarioBros2Env(
34
+ init_config=config,
35
+ render_mode="human", # "human" or None
36
+ action_type="simple" # "simple" (11), "complex" (16), or "all" (256)
37
+ )
38
+
39
+ # Reset environment
40
+ obs, info = env.reset()
41
+
42
+ # Run game loop
43
+ for _ in range(1000):
44
+ action = env.action_space.sample()
45
+ obs, reward, terminated, truncated, info = env.step(action)
46
+
47
+ # Access game state from info dict
48
+ print(f"Lives: {info['life']}, Hearts: {info['hearts']}")
49
+ print(f"Position: ({info['x_pos_global']}, {info['y_pos_global']})")
50
+
51
+ if terminated or truncated:
52
+ obs, info = env.reset()
53
+
54
+ env.close()
55
+ ```
56
+
57
+ ### Initialization Modes
58
+
59
+ ```python
60
+ from smb2_gym.app import InitConfig
61
+
62
+ # 1. Character/Level mode (default)
63
+ config = InitConfig(level="1-1", character="peach")
64
+
65
+ # 2. Built-in ROM variant mode
66
+ config = InitConfig(rom="prg0", save_state="level_1_1.sav")
67
+
68
+ # 3. Custom ROM mode
69
+ config = InitConfig(
70
+ rom_path="/path/to/your/smb2.nes",
71
+ save_state_path="/path/to/save.sav" # Optional
72
+ )
73
+ ```
74
+
75
+ ### Custom Reward Function
76
+
77
+ ```python
78
+ from smb2_gym import SuperMarioBros2Env
79
+ from smb2_gym.app import InitConfig
80
+
81
+ class CustomSMB2Env(SuperMarioBros2Env):
82
+ def step(self, action):
83
+ obs, reward, terminated, truncated, info = super().step(action)
84
+
85
+ # Custom reward based on x-position progress
86
+ reward = info['x_pos_global'] / 100.0
87
+
88
+ # Bonus for collecting cherries
89
+ reward += info['cherries'] * 10
90
+
91
+ # Bonus for hearts
92
+ reward += info['hearts'] * 5
93
+
94
+ # Penalty for losing a life
95
+ if info.get('life_lost'):
96
+ reward -= 100
97
+
98
+ return obs, reward, terminated, truncated, info
99
+
100
+ # Usage
101
+ config = InitConfig(level="1-1", character="luigi")
102
+ env = CustomSMB2Env(init_config=config, action_type="simple")
103
+ ```
104
+
105
+ ## Play as Human
106
+
107
+ The package includes a human-playable interface with multiple initialization modes:
108
+
109
+ ### Character/Level Mode (Default)
110
+ ```bash
111
+ # Play as Luigi on level 1-1
112
+ smb2-play --level 1-1 --char luigi --scale 3
113
+
114
+ # TODO: Yet to implement all levels states
115
+ # Play as Peach on level 2-3
116
+ smb2-play --level 2-3 --char peach
117
+ ```
118
+
119
+ ### Built-in ROM Variant Mode
120
+ ```bash
121
+ # Use specific ROM variant with save state
122
+ smb2-play --rom prg0_edited --save-state easy_combined_curriculum.sav
123
+ ```
124
+
125
+ ### Custom ROM Mode
126
+ ```bash
127
+ # Use your own ROM file
128
+ smb2-play --custom-rom /path/to/smb2.nes
129
+
130
+ # Use custom ROM with save state
131
+ smb2-play --custom-rom /path/to/smb2.nes --custom-state /path/to/save.sav
132
+
133
+ # Start from beginning without save state
134
+ smb2-play --custom-rom /path/to/smb2.nes --no-save-state
135
+ ```
136
+
137
+ ### Controls
138
+
139
+ **Primary Controls:**
140
+ - Arrow Keys: Move
141
+ - Z: A button (Jump)
142
+ - X: B button (Pick up/Throw)
143
+ - Enter: Start
144
+ - Right Shift: Select
145
+ - P: Pause
146
+ - R: Reset
147
+ - ESC: Quit
148
+
149
+ **Alternative Controls:**
150
+ - WASD: Move
151
+ - J: A button
152
+ - K: B button
153
+
154
+ **Save States:**
155
+ - F5: Save state
156
+ - F9: Load state
157
+
158
+ ### CLI Options
159
+
160
+ **Character/Level Mode:**
161
+ - `--level`: Level to play (1-1 through 7-2, default: 1-1)
162
+ - `--char`: Character (mario, luigi, peach, toad, default: luigi)
163
+
164
+ **Built-in ROM Mode:**
165
+ - `--rom`: ROM variant (prg0, prg0_edited)
166
+ - `--save-state`: Save state filename
167
+
168
+ **Custom ROM Mode:**
169
+ - `--custom-rom`: Path to custom ROM file
170
+ - `--custom-state`: Path to custom save state (optional)
171
+ - `--no-save-state`: Start from beginning without loading save state
172
+
173
+ **Display:**
174
+ - `--scale`: Display scale factor (1-4, default: 3)
175
+
176
+ ## Disclaimer
177
+
178
+ This project is for educational and research purposes only. Users must provide their own legally obtained ROM files.
@@ -0,0 +1,86 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61.0", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [tool.setuptools]
6
+ include-package-data = true
7
+
8
+ [tool.setuptools.packages.find]
9
+ include = ["smb2_gym", "smb2_gym.*"]
10
+
11
+ [tool.setuptools.package-data]
12
+ smb2_gym = ["_nes/**/*"]
13
+
14
+ [project]
15
+ name = "smb2-gym"
16
+ version = "0.2.0"
17
+ authors = [{ name = "Alexei Quick" }]
18
+ description = "Super Mario Bros 2 (Europe) Gymnasium Environment"
19
+ readme = "README.md"
20
+ requires-python = ">=3.8"
21
+ license = {text = "MIT"}
22
+ classifiers = [
23
+ "Development Status :: 3 - Alpha",
24
+ "Intended Audience :: Science/Research",
25
+ "Operating System :: OS Independent",
26
+ "Programming Language :: Python :: 3",
27
+ "Programming Language :: Python :: 3.8",
28
+ "Programming Language :: Python :: 3.9",
29
+ "Programming Language :: Python :: 3.10",
30
+ "Programming Language :: Python :: 3.11",
31
+ "Programming Language :: Python :: 3.12",
32
+ "Programming Language :: Python :: 3.13",
33
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
34
+ ]
35
+ dependencies = [
36
+ "gymnasium>=0.29.0",
37
+ "numpy>=1.21.0",
38
+ "pygame>=2.1.0",
39
+ "tetanes-py>=0.1.1",
40
+ ]
41
+
42
+ [project.urls]
43
+ "Homepage" = "https://github.com/alexeiquickcode/smb2-gym"
44
+ "Bug Tracker" = "https://github.com/alexeiquickcode/smb2-gym/issues"
45
+
46
+ [project.optional-dependencies]
47
+ dev = ["pytest", "isort", "yapf", "basedpyright"]
48
+
49
+ [project.scripts]
50
+ smb2-play = "smb2_gym.play_human:main"
51
+
52
+ [tool.yapf]
53
+ based_on_style = "google"
54
+ column_limit = 100
55
+ indent_width = 4
56
+ dedent_closing_brackets = true
57
+ INDENT_DICTIONARY_VALUE = true
58
+ split_before_logical_operator = true
59
+ spaces_before_comment = 2
60
+ split_complex_comprehension = true
61
+ split_arguments_when_comma_terminated = true
62
+ split_all_comma_separated_values = false
63
+ SPLIT_ALL_TOP_LEVEL_COMMA_SEPARATED_VALUES = false
64
+ SPLIT_BEFORE_ARITHMETIC_OPERATOR = true
65
+ SPLIT_BEFORE_DOT = true
66
+
67
+ [tool.isort]
68
+ line_length = 120
69
+ multi_line_output = 3
70
+ include_trailing_comma = true
71
+ force_grid_wrap = 2
72
+ use_parentheses = true
73
+ ensure_newline_before_comments = true
74
+ lines_after_imports = 2
75
+ skip_gitignore = true
76
+ known_first_party = []
77
+
78
+ [tool.pytest.ini_options]
79
+ testpaths = ["tests"]
80
+ python_files = ["test_*.py"]
81
+ python_classes = ["Test*"]
82
+ python_functions = ["test_*"]
83
+ addopts = "-v --tb=short"
84
+ markers = [
85
+ "slow: marks tests as slow (deselect with '-m \"not slow\"')"
86
+ ]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,35 @@
1
+ """Super Mario Bros 2 Gymnasium Environment."""
2
+
3
+ from .actions import (
4
+ COMPLEX_ACTIONS,
5
+ SIMPLE_ACTIONS,
6
+ ActionType,
7
+ action_to_buttons,
8
+ actions_to_buttons,
9
+ buttons_to_action,
10
+ buttons_to_action_index,
11
+ get_action_meanings,
12
+ )
13
+ from .app import (
14
+ create_info_panel,
15
+ draw_info,
16
+ get_required_info_height,
17
+ )
18
+ from .constants import *
19
+ from .smb2_env import SuperMarioBros2Env
20
+
21
+ __version__ = "0.2.0"
22
+ __all__ = [
23
+ "SuperMarioBros2Env",
24
+ "SIMPLE_ACTIONS",
25
+ "COMPLEX_ACTIONS",
26
+ "ActionType",
27
+ "action_to_buttons",
28
+ "buttons_to_action",
29
+ "buttons_to_action_index",
30
+ "actions_to_buttons",
31
+ "get_action_meanings",
32
+ "create_info_panel",
33
+ "get_required_info_height",
34
+ "draw_info",
35
+ ]
@@ -0,0 +1,168 @@
1
+ """Action mappings for Super Mario Bros 2."""
2
+
3
+ from typing import Literal
4
+
5
+ import numpy as np
6
+
7
+ from .constants import (
8
+ BUTTON_A,
9
+ BUTTON_B,
10
+ BUTTON_DOWN,
11
+ BUTTON_LEFT,
12
+ BUTTON_RIGHT,
13
+ BUTTON_SELECT,
14
+ BUTTON_START,
15
+ BUTTON_UP,
16
+ )
17
+
18
+ ActionType = Literal["all", "complex", "simple"]
19
+
20
+ SIMPLE_ACTIONS = [
21
+ ['NOOP'],
22
+ ['right'],
23
+ ['left'],
24
+ ['up'], # Enter door
25
+ ['A'], # Jump
26
+ ['B'], # Pick up/throw
27
+ ['right', 'A'], # Jump right
28
+ ['left', 'A'], # Jump left
29
+ ['right', 'B'], # Pick up right
30
+ ['left', 'B'], # Pick up left
31
+ ['down'], # Crouch/charge super jump (hold 77+ frames then press A)
32
+ ]
33
+
34
+ COMPLEX_ACTIONS = [
35
+ ['NOOP'],
36
+ ['right'],
37
+ ['right', 'A'], # Run right
38
+ ['right', 'B'], # Pick up right
39
+ ['right', 'A', 'B'], # Jump throw right
40
+ ['left'],
41
+ ['left', 'A'], # Run left
42
+ ['left', 'B'], # Pick up left
43
+ ['left', 'A', 'B'], # Jump throw left
44
+ ['down'], # Charge super jump (hold 77+ frames then press A)
45
+ ['down', 'B'], # Pick up from ground
46
+ ['up'], # Enter door
47
+ ['A'], # Jump
48
+ ['B'], # Pick up/throw
49
+ ['A', 'B'], # Jump and throw
50
+ ]
51
+
52
+ BUTTON_MAP = {
53
+ 'a': BUTTON_A,
54
+ 'A': BUTTON_A,
55
+ 'b': BUTTON_B,
56
+ 'B': BUTTON_B,
57
+ 'select': BUTTON_SELECT,
58
+ 'start': BUTTON_START,
59
+ 'up': BUTTON_UP,
60
+ 'down': BUTTON_DOWN,
61
+ 'left': BUTTON_LEFT,
62
+ 'right': BUTTON_RIGHT,
63
+ 'noop': None
64
+ }
65
+
66
+
67
+ def actions_to_buttons(actions: list[str]) -> np.ndarray:
68
+ """Convert action list to button array.
69
+
70
+ Args:
71
+ actions: list of action strings (e.g., ['right', 'A'])
72
+
73
+ Returns:
74
+ Boolean array of shape (8,) for button states
75
+ """
76
+ buttons = np.zeros(8, dtype=bool)
77
+
78
+ for action in actions:
79
+ action_lower = action.lower()
80
+ if action_lower in BUTTON_MAP:
81
+ button_idx = BUTTON_MAP[action_lower]
82
+ if button_idx is not None:
83
+ buttons[button_idx] = True
84
+
85
+ return buttons
86
+
87
+
88
+ def action_to_buttons(action: int) -> np.ndarray:
89
+ """Convert discrete action to button array using all combinations.
90
+
91
+ Args:
92
+ action: Discrete action (0-255)
93
+
94
+ Returns:
95
+ Boolean array of shape (8,) for button states
96
+ """
97
+ # Each bit represents a button
98
+ buttons = np.zeros(8, dtype=bool)
99
+ for i in range(8):
100
+ buttons[i] = bool(action & (1 << i))
101
+ return buttons
102
+
103
+
104
+ def buttons_to_action(buttons: np.ndarray) -> int:
105
+ """Convert button array to discrete action.
106
+
107
+ Args:
108
+ buttons: Boolean array of shape (8,)
109
+
110
+ Returns:
111
+ Discrete action (0-255)
112
+ """
113
+ action = 0
114
+ for i in range(8):
115
+ if buttons[i]:
116
+ action |= (1 << i)
117
+ return action
118
+
119
+
120
+ def buttons_to_action_index(buttons: np.ndarray, action_set: list[list[str]]) -> int:
121
+ """Convert button array to action index for a given action set.
122
+
123
+ Args:
124
+ buttons: Boolean array of shape (8,)
125
+ action_set: List of action combinations (e.g., SIMPLE_ACTIONS, COMPLEX_ACTIONS)
126
+
127
+ Returns:
128
+ Action index in the given action set
129
+ """
130
+ # Convert button state to action list
131
+ pressed = []
132
+ if buttons[BUTTON_RIGHT]:
133
+ pressed.append('right')
134
+ elif buttons[BUTTON_LEFT]:
135
+ pressed.append('left')
136
+ elif buttons[BUTTON_DOWN]:
137
+ pressed.append('down')
138
+ elif buttons[BUTTON_UP]:
139
+ pressed.append('up')
140
+
141
+ if buttons[BUTTON_A]:
142
+ pressed.append('A')
143
+ if buttons[BUTTON_B]:
144
+ pressed.append('B')
145
+
146
+ # Find matching action in the action set
147
+ for i, action_combo in enumerate(action_set):
148
+ if set(pressed) == set([a for a in action_combo if a.upper() != 'NOOP']):
149
+ return i
150
+
151
+ return 0 # NOOP
152
+
153
+
154
+ def get_action_meanings() -> list[list[str]]:
155
+ """Get human-readable meanings for all 256 actions.
156
+
157
+ Returns:
158
+ list of action meanings
159
+ """
160
+ meanings = []
161
+ button_names = ['A', 'B', 'SELECT', 'START', 'UP', 'DOWN', 'LEFT', 'RIGHT']
162
+
163
+ for action in range(256):
164
+ buttons = action_to_buttons(action)
165
+ pressed = [button_names[i] for i in range(8) if buttons[i]]
166
+ meanings.append(pressed if pressed else ['NOOP'])
167
+
168
+ return meanings