Skills › Scientific Agent Skills › stable-baselines3 › custom_env_template.py
custom_env_template.py
325 lines · skills/stable-baselines3/scripts/custom_env_template.py at 5f2bc03
1"""
2Template for creating custom Gymnasium environments compatible with Stable Baselines3.
4This template demonstrates:
5- Proper Gymnasium environment structure
6- Observation and action space definition
7- Step and reset implementation
8- Validation with SB3's env_checker
9- Registration with Gymnasium
10"""
12import gymnasium as gym
13from gymnasium import spaces
14import numpy as np
17class CustomEnv(gym.Env):
18 """
19 Custom Gymnasium Environment Template.
21 This is a template for creating custom environments that work with
22 Stable Baselines3. Modify the observation space, action space, reward
23 function, and state transitions to match your specific problem.
25 Example:
26 A simple grid world where the agent tries to reach a goal position.
27 """
29 # Optional: Provide metadata for rendering modes
30 metadata = {"render_modes": ["human", "rgb_array"], "render_fps": 30}
32 def __init__(self, grid_size=5, render_mode=None):
33 """
34 Initialize the environment.
36 Args:
37 grid_size: Size of the grid world (grid_size x grid_size)
38 render_mode: How to render ('human', 'rgb_array', or None)
39 """
40 super().__init__()
42 if not isinstance(grid_size, (int, np.integer)) or isinstance(grid_size, bool) or grid_size < 2:
43 raise ValueError("grid_size must be an integer >= 2")
44 if render_mode not in (None, *self.metadata["render_modes"]):
45 raise ValueError("unsupported render_mode")
46 self.grid_size = grid_size
47 self.render_mode = render_mode
49 # Define action space
50 # Example: 4 discrete actions (up, down, left, right)
51 self.action_space = spaces.Discrete(4)
53 # Define observation space
54 # Agent and goal: [agent_row, agent_col, goal_row, goal_col].
55 # The random goal must be observable for this Markov task.
56 # Note: Use np.float32 for observations (SB3 recommendation)
57 self.observation_space = spaces.Box(
58 low=0,
59 high=grid_size - 1,
60 shape=(4,),
61 dtype=np.float32,
62 )
64 # Alternative observation spaces:
65 # 1. Discrete: spaces.Discrete(n)
66 # 2. Multi-discrete: spaces.MultiDiscrete([n1, n2, ...])
67 # 3. Multi-binary: spaces.MultiBinary(n)
68 # 4. Box (continuous): spaces.Box(low=, high=, shape=, dtype=np.float32)
69 # 5. Dict: spaces.Dict({"key1": space1, "key2": space2})
71 # For image observations (e.g., 84x84 RGB image):
72 # self.observation_space = spaces.Box(
73 # low=0,
74 # high=255,
75 # shape=(3, 84, 84), # (channels, height, width) - channel-first
76 # dtype=np.uint8,
77 # )
79 # Initialize state
80 self._agent_position = None
81 self._goal_position = None
83 def reset(self, seed=None, options=None):
84 """
85 Reset the environment to initial state.
87 Args:
88 seed: Random seed for reproducibility
89 options: Additional options (optional)
91 Returns:
92 observation: Initial observation
93 info: Additional information dictionary
94 """
95 # Set seed for reproducibility
96 super().reset(seed=seed)
98 # Initialize agent position randomly
99 self._agent_position = self.np_random.integers(0, self.grid_size, size=2)
101 # Initialize goal position (different from agent)
102 self._goal_position = self.np_random.integers(0, self.grid_size, size=2)
103 while np.array_equal(self._agent_position, self._goal_position):
104 self._goal_position = self.np_random.integers(0, self.grid_size, size=2)
106 observation = self._get_obs()
107 info = self._get_info()
109 return observation, info
111 def step(self, action):
112 """
113 Execute one step in the environment.
115 Args:
116 action: Action to take
118 Returns:
119 observation: New observation
120 reward: Reward for this step
121 terminated: Whether episode has ended (goal reached)
122 truncated: Whether episode was truncated (time limit, etc.)
123 info: Additional information dictionary
124 """
125 if self._agent_position is None:
126 raise RuntimeError("reset() must be called before step()")
127 action = np.asarray(action)
128 if action.shape != () or not self.action_space.contains(action):
129 raise ValueError("action must be one discrete value in [0, 3]")
130 # Map action to direction (0: up, 1: down, 2: left, 3: right)
131 direction = np.array([
132 [-1, 0], # up
133 [1, 0], # down
134 [0, -1], # left
135 [0, 1], # right
136 ])[int(action)]
138 # Update agent position (clip to stay within grid)
139 self._agent_position = np.clip(
140 self._agent_position + direction,
141 0,
142 self.grid_size - 1,
143 )
145 # Check if goal is reached
146 terminated = np.array_equal(self._agent_position, self._goal_position)
148 # Calculate reward
149 if terminated:
150 reward = 1.0 # Goal reached
151 else:
152 # Negative reward based on distance to goal (encourages efficiency)
153 distance = np.linalg.norm(self._agent_position - self._goal_position)
154 reward = -0.1 * distance
156 # Episode not truncated in this example (no time limit)
157 truncated = False
159 observation = self._get_obs()
160 info = self._get_info()
162 return observation, reward, terminated, truncated, info
164 def _get_obs(self):
165 """
166 Get current observation.
168 Returns:
169 observation: Current state as defined by observation_space
170 """
171 # Return a fresh array including the observable goal.
172 return np.concatenate((self._agent_position, self._goal_position)).astype(np.float32)
174 # For dict observations:
175 # return {
176 # "agent": self._agent_position.astype(np.float32),
177 # "goal": self._goal_position.astype(np.float32),
178 # }
180 def _get_info(self):
181 """
182 Get additional information (for debugging/logging).
184 Returns:
185 info: Dictionary with additional information
186 """
187 return {
188 "agent_position": self._agent_position.copy(),
189 "goal_position": self._goal_position.copy(),
190 "distance_to_goal": np.linalg.norm(
191 self._agent_position - self._goal_position
192 ),
193 }
195 def render(self):
196 """
197 Render the environment.
199 Returns:
200 Rendered frame (if render_mode is 'rgb_array')
201 """
202 if self.render_mode == "human":
203 # Print simple text-based rendering
204 grid = np.zeros((self.grid_size, self.grid_size), dtype=str)
205 grid[:, :] = "."
206 grid[tuple(self._agent_position)] = "A"
207 grid[tuple(self._goal_position)] = "G"
209 print("\n" + "=" * (self.grid_size * 2 + 1))
210 for row in grid:
211 print(" ".join(row))
212 print("=" * (self.grid_size * 2 + 1) + "\n")
214 elif self.render_mode == "rgb_array":
215 # Return RGB array for video recording
216 canvas = np.zeros((
217 self.grid_size * 50,
218 self.grid_size * 50,
219 3
220 ), dtype=np.uint8)
221 for position, color in ((self._goal_position, (0, 200, 0)),
222 (self._agent_position, (40, 100, 255))):
223 row, col = position * 50
224 canvas[row:row + 50, col:col + 50] = color
225 return canvas
227 def close(self):
228 """
229 Clean up environment resources.
230 """
231 pass
234# Optional: Register the environment with Gymnasium
235# This allows creating the environment with gym.make("CustomEnv-v0")
236gym.register(
237 id="CustomEnv-v0",
238 entry_point=__name__ + ":CustomEnv",
239 max_episode_steps=100,
240)
243def validate_environment():
244 """
245 Validate the custom environment with SB3's env_checker.
246 """
247 from stable_baselines3.common.env_checker import check_env
249 print("Validating custom environment...")
250 env = CustomEnv()
251 check_env(env, warn=True)
252 print("Environment validation passed!")
255def test_environment():
256 """
257 Test the custom environment with random actions.
258 """
259 print("Testing environment with random actions...")
260 env = CustomEnv(render_mode="human")
262 obs, info = env.reset()
263 print(f"Initial observation: {obs}")
264 print(f"Initial info: {info}")
266 for step in range(10):
267 action = env.action_space.sample() # Random action
268 obs, reward, terminated, truncated, info = env.step(action)
270 print(f"\nStep {step + 1}:")
271 print(f" Action: {action}")
272 print(f" Observation: {obs}")
273 print(f" Reward: {reward:.3f}")
274 print(f" Terminated: {terminated}")
275 print(f" Info: {info}")
277 env.render()
279 if terminated or truncated:
280 print("Episode finished!")
281 break
283 env.close()
286def train_on_custom_env():
287 """
288 Train a PPO agent on the custom environment.
289 """
290 from stable_baselines3 import PPO
292 print("Training PPO agent on custom environment...")
294 # Create environment
295 env = gym.make("CustomEnv-v0")
297 # Validate first
298 from stable_baselines3.common.env_checker import check_env
299 check_env(env, warn=True)
301 # Train agent
302 model = PPO("MlpPolicy", env, verbose=1)
303 model.learn(total_timesteps=10000)
305 # Test trained agent
306 obs, info = env.reset()
307 for _ in range(20):
308 action, _states = model.predict(obs, deterministic=True)
309 obs, reward, terminated, truncated, info = env.step(action)
310 if terminated or truncated:
311 print(f"Episode ended: goal={terminated}, timeout={truncated}, reward={reward}")
312 break
314 env.close()
317if __name__ == "__main__":
318 # Validate the environment
319 validate_environment()
321 # Test with random actions
322 # test_environment()
324 # Train an agent
325 # train_on_custom_env()