update frozen lake

Author: jakegrigsby

amago/envs/builtin/toy_gym.py

@@ -36,6 +36,14 @@ class MetaFrozenLake(gym.Env):
         recover_mode: If False, falling through the ice terminates the
             episode. If True, the agent is allowed to recover to its
             previous position but receives a penalty. Defaults to False.
+        max_episode_steps: Maximum steps per attempt before a forced
+            soft reset. Defaults to N² (standard) or 2*N² (hard).
+        show_k_progress: If True, include current_k / k_episodes in
+            observations. Set to False to hide trial progress (useful
+            for testing length extrapolation with different k values).
+            Defaults to True.
+        slip_chance: Probability that a movement action is replaced by
+            a no-op (agent stays in place). Defaults to 0.0.
     """
 
     def __init__(
@@ -44,13 +52,24 @@ def __init__(
         k_episodes: int = 10,
         hard_mode: bool = False,
         recover_mode: bool = False,
+        max_episode_steps: int | None = None,
+        show_k_progress: bool = True,
+        slip_chance: float = 0.0,
     ):
         self.size = size
         self.k_episodes = k_episodes
         self.action_space = gym.spaces.Discrete(5)
         self.observation_space = gym.spaces.Box(shape=(4,), low=0.0, high=1.0)
         self.hard_mode = hard_mode
         self.recover_mode = recover_mode
+        base_steps = size * size * (2 if hard_mode else 1)
+        if slip_chance > 0:
+            base_steps = int(base_steps / (1.0 - slip_chance))
+        self.max_episode_steps = (
+            max_episode_steps if max_episode_steps is not None else base_steps
+        )
+        self.show_k_progress = show_k_progress
+        self.slip_chance = slip_chance
         self.reset()
 
     def reset(self, *args, **kwargs):
@@ -69,30 +88,30 @@ def make_obs(self, reset_signal: bool):
             y = min(max(self.y + random.choice([-1, 0, 1]), 0), self.size - 1)
         else:
             x, y = self.x, self.y
+        k_obs = self.current_k / self.k_episodes if self.show_k_progress else 0.0
         return np.array(
-            [
-                x / self.size,
-                y / self.size,
-                reset_signal,
-                self.current_k / self.k_episodes,
-            ],
+            [x / self.size, y / self.size, reset_signal, k_obs],
             dtype=np.float32,
         )
 
     def soft_reset(self):
         self.active_map = copy.deepcopy(self.current_map)
         self.x, self.y = 0, 0
+        self.episode_steps = 0
         obs = self.make_obs(reset_signal=True)
         return obs, {}
 
     def step(self, action):
         assert self.action_space.contains(action)
+        self.episode_steps += 1
+        if self.slip_chance > 0 and action != 0 and random.random() < self.slip_chance:
+            action = 0
         move_x, move_y = self.action_mapping[action]
         next_x = max(min(self.x + move_x, self.size - 1), 0)
         next_y = max(min(self.y + move_y, self.size - 1), 0)
 
         if (
-            (self.x, self.y) != (next_y, next_y)
+            (self.x, self.y) != (next_x, next_y)
             and self.hard_mode
             and random.random() < 0.33
         ):
@@ -115,6 +134,14 @@ def step(self, action):
         self.x = next_x
         self.y = next_y
 
+        timed_out = (
+            not soft_reset
+            and self.max_episode_steps is not None
+            and self.episode_steps >= self.max_episode_steps
+        )
+        if timed_out:
+            soft_reset = True
+
         if soft_reset:
             next_state, info = self.soft_reset()
             success = on == "G"

examples/00_meta_frozen_lake.py

@@ -19,11 +19,27 @@ def add_cli(parser):
     parser.add_argument("--log", action="store_true")
     parser.add_argument("--trials", type=int, default=1)
     parser.add_argument("--lake_size", type=int, default=5)
-    parser.add_argument("--k_episodes", type=int, default=15)
+    parser.add_argument("--k_episodes", type=int, default=10)
     parser.add_argument("--hard_mode", action="store_true")
     parser.add_argument("--recover_mode", action="store_true")
-    parser.add_argument("--max_rollout_length", type=int, default=512)
-    parser.add_argument("--max_seq_len", type=int, default=512)
+    parser.add_argument("--slip_chance", type=float, default=0.0)
+    parser.add_argument(
+        "--max_episode_steps",
+        type=int,
+        default=None,
+        help="Max steps per attempt. Default: N² (standard) or 2*N² (hard).",
+    )
+    parser.add_argument(
+        "--hide_k_progress",
+        action="store_true",
+        help="Hide current_k/k_episodes from observations (for length extrapolation tests).",
+    )
+    parser.add_argument(
+        "--max_seq_len",
+        type=int,
+        default=None,
+        help="Training sequence length. Default: max_episode_steps * k_episodes (full trajectory).",
+    )
     return parser
 
 
@@ -35,6 +51,19 @@ def add_cli(parser):
     if args.log:
         import wandb
 
+    lake_kwargs = dict(
+        size=args.lake_size,
+        k_episodes=args.k_episodes,
+        hard_mode=args.hard_mode,
+        recover_mode=args.recover_mode,
+        max_episode_steps=args.max_episode_steps,
+        show_k_progress=not args.hide_k_progress,
+        slip_chance=args.slip_chance,
+    )
+    max_ep_steps = MetaFrozenLake(**lake_kwargs).max_episode_steps
+    max_rollout_length = max_ep_steps * args.k_episodes
+    max_seq_len = args.max_seq_len or max_rollout_length
+
     config = {}
     # configure trajectory encoder (seq2seq memory model)
     traj_encoder_type = cli_utils.switch_traj_encoder(
@@ -47,7 +76,6 @@ def add_cli(parser):
     tstep_encoder_type = cli_utils.switch_tstep_encoder(
         config, arch="ff", n_layers=1, d_hidden=128, d_output=64, normalize_inputs=False
     )
-
     # we're using the default exploration strategy but being overly verbose about it for the example
     exploration_wrapper_type = cli_utils.switch_exploration(
         config,
@@ -70,28 +98,21 @@ def add_cli(parser):
         )
         # save checkpoints alongside the buffer
         ckpt_dir = args.buffer_dir
-
         # wrap environment
         make_env = lambda: AMAGOEnv(
-            MetaFrozenLake(
-                k_episodes=args.k_episodes,
-                size=args.lake_size,
-                hard_mode=args.hard_mode,
-                recover_mode=args.recover_mode,
-            ),
+            MetaFrozenLake(**lake_kwargs),
             env_name=f"meta_frozen_lake_k{args.k_episodes}_{args.lake_size}x{args.lake_size}"
             + ("_hard" if args.hard_mode else "_easy")
             + ("_recover" if args.recover_mode else "_reset"),
         )
 
-        # create `Experiment`
         experiment = amago.Experiment(
             make_train_env=make_env,
             make_val_env=make_env,
-            max_seq_len=args.max_seq_len,
-            traj_save_len=args.max_rollout_length,
+            max_seq_len=max_seq_len,
+            traj_save_len=max_rollout_length,
             dataset=dset,
-            ckpt_base_dir=ckpt_dir,
+            ckpt_base_dir=args.buffer_dir,
             agent_type=agent_type,
             exploration_wrapper_type=exploration_wrapper_type,
             tstep_encoder_type=tstep_encoder_type,
@@ -102,18 +123,17 @@ def add_cli(parser):
             wandb_group_name=group_name,
             epochs=700 if not args.hard_mode else 900,
             parallel_actors=32,
-            train_timesteps_per_epoch=args.max_rollout_length,
+            train_timesteps_per_epoch=max_rollout_length,
             train_batches_per_epoch=1000,
             val_interval=20,
-            val_timesteps_per_epoch=args.max_rollout_length * 2,
+            val_timesteps_per_epoch=max_rollout_length * 2,
             ckpt_interval=200,
             env_mode="sync",
         )
 
-        # start experiment (build envs, policies, etc.)
         experiment.start()
-        # run training
         experiment.learn()
         experiment.evaluate_test(make_env, timesteps=10_000)
         experiment.delete_buffer_from_disk()
-        wandb.finish()
+        if args.log:
+            wandb.finish()