Merge pull request #103 from UT-Austin-RPL/zs

Env Updates

Author: jakegrigsby

amago/agent.py

@@ -763,7 +763,7 @@ def get_actions(
                 actions = action_dists.mean
         # get intended gamma distribution (always in -1 idx)
         actions = actions[..., -1, :]
-        dtype = torch.uint8 if (self.discrete or self.multibinary) else torch.float32
+        dtype = torch.int64 if (self.discrete or self.multibinary) else torch.float32
         return actions.to(dtype=dtype), hidden_state
 
     @torch.no_grad()

amago/envs/builtin/half_cheetah_v4_vel.py

@@ -3,7 +3,9 @@
 """
 
 import random
+from typing import List
 
+import gymnasium as gym
 from gymnasium.envs.mujoco.half_cheetah_v4 import HalfCheetahEnv
 import numpy as np
 
@@ -37,7 +39,19 @@ def step(self, action):
 
 
 class HalfCheetahV4LogVelocity(_HalfCheetahV4ExposeVelReward):
-    """A wrapper around HalfCheetah-V4 that will automatically log velocity metrics when used in AMAGO."""
+    """A wrapper around HalfCheetah-V4 that will automatically log velocity metrics when used in AMAGO.
+
+    Args:
+        max_episode_steps: Step horizon at which the inner episode is
+            truncated. Default 1000 matches the original gym registration.
+            Lower values are useful for k-episode meta-trial wrappers that
+            want to fit multiple inner episodes inside a fixed sequence
+            length.
+    """
+
+    def __init__(self, max_episode_steps: int = 1000, **kwargs):
+        super().__init__(**kwargs)
+        self.max_episode_steps = int(max_episode_steps)
 
     def reset(self, *args, **kwargs):
         obs, info = super().reset(*args, **kwargs)
@@ -51,7 +65,7 @@ def step(self, action):
         # gym registration that would normally handle it.
         self.step_count += 1
         self._velocity_history.append(info["x_velocity"])
-        truncated = truncated or self.step_count >= 1000
+        truncated = truncated or self.step_count >= self.max_episode_steps
         if terminated or truncated:
             v = np.array(self._velocity_history)
             info[f"{AMAGO_ENV_LOG_PREFIX} Mean Velocity"] = v.mean().item()
@@ -69,6 +83,17 @@ class HalfCheetahV4_MetaVelocity(HalfCheetahV4LogVelocity):
     Reward terms are based on the version featured in the VariBAD codebase
     (https://github.com/lmzintgraf/varibad/blob/57e1795be142ace52d0c353097acf193d9067200/environments/mujoco/half_cheetah_vel.py#L8)
 
+    A single hidden ``target_velocity`` persists across ``k_episodes`` inner
+    episodes (each truncated by the parent class at ``max_episode_steps``).
+    Between inner episodes the simulator is soft-reset; the hidden task is
+    only resampled when ``reset(new_task=True)`` is called from outside (the
+    default at the start of every meta-trial). A soft-reset flag (0/1) is
+    appended to the obs so the agent can detect inner-episode boundaries
+    without losing task identity. Per-episode metrics from the parent class
+    (Mean / Max / Avg. Last 50 Timestep Velocity) are promoted to
+    ``Trial {i} ...`` keys, matching the meta-RL logging convention used by
+    ``KEpisodeMetaworld``.
+
     Args:
         ctrl_cost_weight: Defaults to half the normal ctrl cost, as in the original HalfCheetahVelocity task.
         velocity_reward_weight: Defaults to 1.0.
@@ -80,6 +105,10 @@ class HalfCheetahV4_MetaVelocity(HalfCheetahV4LogVelocity):
 
         task_min_velocity: Minimum target velocity that determines the hidden task. Defaults to 0.0.
         task_max_velocity: Maximum target velocity that determines the hidden task. Defaults to 3.0.
+        k_episodes: Number of inner episodes per meta-trial. Default 1
+            recovers the unwrapped single-episode task. Set ``>=2`` for
+            true meta-RL trials (e.g. 3 inner episodes of 200 steps with
+            ``max_episode_steps=200``).
     """
 
     def __init__(
@@ -93,6 +122,7 @@ def __init__(
         # We can use the HalfCheetahV4LogVelocity env above to verify this.
         task_min_velocity: float = 0.0,
         task_max_velocity: float = 3.0,
+        k_episodes: int = 1,
         **kwargs,
     ):
         super().__init__(
@@ -102,8 +132,24 @@ def __init__(
         )
         self.task_min_velocity = task_min_velocity
         self.task_max_velocity = task_max_velocity
+        self.k_episodes = int(k_episodes)
+
+        # Append a soft-reset flag (0/1) to the parent observation space.
+        flat = self.observation_space
+        low = np.append(flat.low.astype(np.float32), np.float32(0.0))
+        high = np.append(flat.high.astype(np.float32), np.float32(1.0))
+        self.observation_space = gym.spaces.Box(low=low, high=high, dtype=np.float32)
+
+        self.target_velocity: float = 0.0
+        self.current_trial: int = 0
+        self._trial_returns: List[float] = []
+        self._current_trial_return: float = 0.0
         self.reset()
 
+    @staticmethod
+    def _augment(obs: np.ndarray, soft_reset: bool) -> np.ndarray:
+        return np.append(obs, float(soft_reset)).astype(np.float32)
+
     def velocity_reward_term(self, x_velocity):
         return (
             self._forward_reward_weight
@@ -117,16 +163,64 @@ def sample_target_velocity(self):
         # tasks are different across async parallel actors!
         return random.uniform(self.task_min_velocity, self.task_max_velocity)
 
-    def reset(self, *args, **kwargs):
-        obs, info = super().reset(*args, **kwargs)
-        self.target_velocity = self.sample_target_velocity()
-        return obs, info
+    def reset(self, *, new_task: bool = True, seed=None, options=None):
+        # ``new_task=False`` is reserved for soft resets *inside* step();
+        # external callers should use the default to start a fresh meta-trial.
+        obs, info = super().reset(seed=seed, options=options)
+        if new_task:
+            self.target_velocity = self.sample_target_velocity()
+            self.current_trial = 0
+            self._trial_returns = []
+        self._current_trial_return = 0.0
+        return self._augment(obs, soft_reset=True), info
+
+    def _log_per_trial_metrics(self, info: dict) -> None:
+        ti = self.current_trial
+        for suffix in (
+            "Mean Velocity",
+            "Max Velocity",
+            "Avg. Last 50 Timestep Velocity",
+        ):
+            src = f"{AMAGO_ENV_LOG_PREFIX} {suffix}"
+            if src in info:
+                info[f"{AMAGO_ENV_LOG_PREFIX} Trial {ti} {suffix}"] = info.pop(src)
+        achieved_key = (
+            f"{AMAGO_ENV_LOG_PREFIX} Trial {ti} Avg. Last 50 Timestep Velocity"
+        )
+        if achieved_key in info:
+            info[
+                f"{AMAGO_ENV_LOG_PREFIX} Trial {ti} Target Velocity Error at Final 50 Timesteps"
+            ] = abs(self.target_velocity - info[achieved_key])
+        info[f"{AMAGO_ENV_LOG_PREFIX} Trial {ti} Return"] = self._current_trial_return
 
     def step(self, action):
         obs, reward, terminated, truncated, info = super().step(action)
+        self._current_trial_return += float(reward)
+        soft_reset = False
+
         if terminated or truncated:
-            achieved = info[f"{AMAGO_ENV_LOG_PREFIX} Avg. Last 50 Timestep Velocity"]
-            info[
-                f"{AMAGO_ENV_LOG_PREFIX} Target Velocity Error at Final 50 Timesteps"
-            ] = abs(self.target_velocity - achieved)
-        return obs, reward, terminated, truncated, info
+            self._log_per_trial_metrics(info)
+            self._trial_returns.append(self._current_trial_return)
+            self.current_trial += 1
+
+            # Always soft-reset the simulator on inner-episode end (mirrors
+            # KEpisodeMetaworld). The augmented obs's soft-reset flag is set
+            # to 1 so the agent can detect the boundary.
+            obs, _ = super().reset()
+            soft_reset = True
+            self._current_trial_return = 0.0
+
+            if self.current_trial >= self.k_episodes:
+                # Meta-trial end is a time-limit truncation, not a terminal
+                # state: the underlying MDP could continue if we kept
+                # observing. Leave terminated=False so the critic
+                # bootstraps as usual.
+                terminated = False
+                truncated = True
+                info[f"{AMAGO_ENV_LOG_PREFIX} Total Return"] = float(
+                    np.sum(self._trial_returns)
+                )
+            else:
+                terminated = truncated = False
+
+        return self._augment(obs, soft_reset), reward, terminated, truncated, info

amago/envs/builtin/metaworld_ml.py

@@ -33,7 +33,13 @@ class Metaworld(AMAGOEnv):
         k_episodes: The number of attempts the policy has to adapt to the current task.
     """
 
-    def __init__(self, benchmark_name: str, split: str, k_episodes: int):
+    def __init__(
+        self,
+        benchmark_name: str,
+        split: str,
+        k_episodes: int,
+        max_episode_length: int = 500,
+    ):
         if benchmark_name == "ml10":
             benchmark = metaworld.ML10()
         elif benchmark_name == "ml45":
@@ -44,7 +50,7 @@ def __init__(self, benchmark_name: str, split: str, k_episodes: int):
             except:
                 raise ValueError(f"Unrecognized metaworld benchmark {benchmark_name}")
 
-        env = KEpisodeMetaworld(benchmark, split, k_episodes)
+        env = KEpisodeMetaworld(benchmark, split, k_episodes, max_episode_length)
         super().__init__(
             env=env,
             env_name=f"metaworld_{benchmark_name}",

amago/envs/env_utils.py

@@ -141,7 +141,7 @@ def action(self, action):
         if isinstance(action, int):
             return action
         elif isinstance(action, np.ndarray) and action.ndim == 1:
-            assert action.shape[0] == 1 and action.dtype == np.uint8
+            assert action.shape[0] == 1
             return action[0].item()
         elif isinstance(action, np.ndarray) and action.ndim == 2:
             assert action.shape[1] == 1

amago/envs/exploration.py

@@ -205,7 +205,7 @@ def add_exploration_noise(self, action: np.ndarray, local_step: np.ndarray):
             use_random = np.expand_dims(self.rng.random(self.batched_envs) <= noise, 1)
             expl_action = (
                 use_random * random_action + (1 - use_random) * action
-            ).astype(np.uint8)
+            ).astype(np.int64)
             assert expl_action.shape == (self.batched_envs, 1)
         else:
             # random noise (TD3-style)

amago/experiment.py

@@ -77,7 +77,7 @@ class Experiment:
     :param wandb_group_name: Group runs on wandb dashboard. **Default:** None.
     :param verbose: Print tqdm bars and info to console. **Default:** True.
     :param log_interval: Log extra metrics every N batches. **Default:** 300.
-    :param padded_sampling: Padding for sampling training subsequences. "none", "left", "right", "both". **Default:** "none".
+    :param padded_sampling: Padding for sampling training subsequences. "none", "left", "right", "both", "center". **Default:** "none".
     :param dloader_workers: Number of DataLoader workers for disk loading. Increase for compressed/large trajs.
 
     .. note::
@@ -522,6 +522,15 @@ def policy(self) -> BaseAgent:
         """Returns the current Agent policy free from the accelerator wrapper."""
         return self.accelerator.unwrap_model(self.policy_aclr)
 
+    def on_interact_step(self, obs: dict) -> None:
+        """Hook called before each action during ``self.interact(...)``.
+
+        Override in a subclass to mutate agent state between steps
+        (e.g. resample ``agent._inference_alpha`` at episode boundaries by
+        inspecting ``obs["episode_id"]``).  No-op in the base implementation.
+        """
+        pass
+
     def interact(
         self,
         envs,
@@ -603,6 +612,7 @@ def get_t():
         obs, rl2s, time_idxs = get_t()
         episodes_finished = 0
         for _ in iter_:
+            self.on_interact_step(obs)
             with torch.no_grad():
                 with self.caster():
                     actions, hidden_state = policy.get_actions(

amago/loading.py

@@ -90,6 +90,9 @@ def random_slice(self, length: int, padded_sampling: str = "none"):
                 where the start of the training sequence is not always the first timestep of the trajectory.
                 "right" --> sample while effectively padding the right side of the sequence for cases
                 where the end of the trajectory may be undersampled.
+                "center" --> uniform random window center, clamped to valid range.
+                Always returns a full-length (never padded) window. Fixes endpoint
+                undersampling of "none" when the trajectory is longer than `length`.
         """
         if len(self) <= length:
             start = 0
@@ -101,6 +104,9 @@ def random_slice(self, length: int, padded_sampling: str = "none"):
             start = self._safe_randrange(-length + 1, len(self) - length + 1)
         elif padded_sampling == "right":
             start = self._safe_randrange(0, len(self) - 1)
+        elif padded_sampling == "center":
+            center = self._safe_randrange(0, len(self))
+            start = min(max(center - length // 2, 0), len(self) - length)
         else:
             raise ValueError(
                 f"Unrecognized `padded_sampling` mode: `{padded_sampling}`"

examples/07_metaworld.py

@@ -38,6 +38,8 @@ def add_cli(parser):
         "amago.nets.actor_critic.NCriticsTwoHot.min_return": -100.0,
         "amago.nets.actor_critic.NCriticsTwoHot.max_return": 5000 * args.k,
         "amago.nets.actor_critic.NCriticsTwoHot.output_bins": 96,
+        #"amago.nets.traj_encoders.TformerTrajEncoder.pos_emb": "rope",
+        "amago.nets.actor_critic.NCriticsTwoHot.d_hidden": 300,
     }
     traj_encoder_type = cli_utils.switch_traj_encoder(
         config,
@@ -49,12 +51,16 @@ def add_cli(parser):
         config, args.agent_type, reward_multiplier=1.0, num_critics=4
     )
     exploration_type = cli_utils.switch_exploration(
-        config, "bilevel", steps_anneal=2_000_000, rollout_horizon=args.k * 500
+        config, "bilevel", steps_anneal=2_000_000, rollout_horizon=args.k * 300
     )
     cli_utils.use_config(config, args.configs)
 
-    make_train_env = lambda: Metaworld(args.benchmark, "train", k_episodes=args.k)
-    make_test_env = lambda: Metaworld(args.benchmark, "test", k_episodes=args.k)
+    make_train_env = lambda: Metaworld(
+        args.benchmark, "train", k_episodes=args.k, max_episode_length=300
+    )
+    make_test_env = lambda: Metaworld(
+        args.benchmark, "test", k_episodes=args.k, max_episode_length=300
+    )
 
     group_name = (
         f"{args.run_name}_metaworld_{args.benchmark}_K_{args.k}_L_{args.max_seq_len}"
@@ -66,13 +72,13 @@ def add_cli(parser):
             make_train_env=make_train_env,
             make_val_env=make_train_env,
             max_seq_len=args.max_seq_len,
-            traj_save_len=min(500 * args.k + 1, args.max_seq_len * 4),
+            traj_save_len=min(300 * args.k + 1, args.max_seq_len * 4),
             group_name=group_name,
             run_name=run_name,
             tstep_encoder_type=FFTstepEncoder,
             traj_encoder_type=traj_encoder_type,
             agent_type=agent_type,
-            val_timesteps_per_epoch=15 * args.k * 500 + 1,
+            val_timesteps_per_epoch=15 * args.k * 300 + 1,
             learning_rate=5e-4,
             grad_clip=2.0,
             exploration_wrapper_type=exploration_type,