diff --git a/frozenlakeplus/README.md b/frozenlakeplus/README.md new file mode 100644 index 0000000000..abd6422092 --- /dev/null +++ b/frozenlakeplus/README.md @@ -0,0 +1,13 @@ +# FrozenLakePlus: Dynamic Ice Tile Environment ❄️ + +This is a custom extension of OpenAI Gym's FrozenLake environment, where the map's slipperiness dynamically changes every few steps. + +## Features + +- Switches between slippery and non-slippery +- Simple to plug into SB3 or custom agents +- Great for teaching non-stationary RL + +## How to Run + +python test_frozenlake_plus.py diff --git a/frozenlakeplus/eval_agent.py b/frozenlakeplus/eval_agent.py new file mode 100644 index 0000000000..524db69dc7 --- /dev/null +++ b/frozenlakeplus/eval_agent.py @@ -0,0 +1,23 @@ +from frozenlake_plus_env import FrozenLakePlus +from stable_baselines3 import PPO +import numpy as np + +env = FrozenLakePlus(dynamic_slippery=False) +model = PPO.load("ppo_frozenlakeplus", env=env) + +n_episodes = 10 +for episode in range(n_episodes): + obs, info = env.reset() + done = False + total_reward = 0 + while not done: + action, _ = model.predict(obs, deterministic=True) + # Convert action to int if it's a numpy array + if isinstance(action, np.ndarray): + action = int(action.item()) + obs, reward, terminated, truncated, info = env.step(action) + total_reward += reward + done = terminated or truncated + print(f"Episode {episode+1}: Total Reward = {total_reward}") + +env.close() \ No newline at end of file diff --git a/frozenlakeplus/frozenlake_plus_env.py b/frozenlakeplus/frozenlake_plus_env.py new file mode 100644 index 0000000000..ffd55b3a3e --- /dev/null +++ b/frozenlakeplus/frozenlake_plus_env.py @@ -0,0 +1,31 @@ +import gymnasium as gym +from gymnasium.envs.toy_text.frozen_lake import FrozenLakeEnv +import numpy as np + +class FrozenLakePlus(FrozenLakeEnv): + def __init__(self, map_name="4x4", is_slippery=True, dynamic_slippery=False, slippery_change_freq=10): + super().__init__(map_name=map_name, is_slippery=is_slippery) + self.dynamic_slippery = dynamic_slippery + self.slippery_change_freq = slippery_change_freq + self.step_count = 0 + self.slippery = is_slippery + + def step(self, action): + self.step_count += 1 + + if self.dynamic_slippery and self.step_count % self.slippery_change_freq == 0: + self.slippery = not self.slippery + self.is_slippery = self.slippery + + obs, reward, terminated, truncated, info = super().step(action) + + # Convert the observation to a numpy array + obs = np.array(obs) + + return obs, reward, terminated, truncated, info + + def reset(self, seed=None, options=None): + self.step_count = 0 + return super().reset(seed=seed, options=options) + +__all__ = ["FrozenLakePlus"] diff --git a/frozenlakeplus/ppo_frozenlakeplus.zip b/frozenlakeplus/ppo_frozenlakeplus.zip new file mode 100644 index 0000000000..25c6946673 Binary files /dev/null and b/frozenlakeplus/ppo_frozenlakeplus.zip differ diff --git a/frozenlakeplus/test_frozenlake_plus.py b/frozenlakeplus/test_frozenlake_plus.py new file mode 100644 index 0000000000..02e482a162 --- /dev/null +++ b/frozenlakeplus/test_frozenlake_plus.py @@ -0,0 +1,13 @@ +from frozenlake_plus_env import FrozenLakePlus + +env = FrozenLakePlus(dynamic_slippery=True, slippery_change_freq=3) +obs, info = env.reset() + +done = False +step = 0 # Initialize step counter +while not done: + obs, reward, terminated, truncated, info = env.step(env.action_space.sample()) + step += 1 # Increment step counter + print(f"Step: {step}, Obs: {obs}, Reward: {reward}, Terminated: {terminated}, Truncated: {truncated}") + done = terminated or truncated +env.close() diff --git a/frozenlakeplus/train_agent.py b/frozenlakeplus/train_agent.py new file mode 100644 index 0000000000..fdcc206625 --- /dev/null +++ b/frozenlakeplus/train_agent.py @@ -0,0 +1,12 @@ +from frozenlake_plus_env import FrozenLakePlus +from stable_baselines3 import PPO +from stable_baselines3.common.env_util import make_vec_env + +env = make_vec_env(lambda: FrozenLakePlus(dynamic_slippery=True), n_envs=1) + +print("Starting training...") +model = PPO("MlpPolicy", env, verbose=2) +model.learn(total_timesteps=10_000) +print("Training finished. Model saved as ppo_frozenlakeplus.zip") + +model.save("ppo_frozenlakeplus") \ No newline at end of file