Isaac-GR00T/gr00t/eval/sim/wrapper/multistep_wrapper.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import defaultdict, deque
from enum import Enum
import warnings

import gymnasium as gym
from gymnasium import spaces
import numpy as np


class AggregateMethod(str, Enum):
    """Supported strategies for reducing a sequence of per-step values.

    Subclassing ``str`` keeps the members interchangeable with their
    string values, so existing callers (and configs) that pass e.g.
    ``"max"`` keep working while the code itself references the typed
    members instead of magic strings.
    """

    MAX = "max"
    MIN = "min"
    MEAN = "mean"
    SUM = "sum"


def stack_repeated(x, n, loc):
    return np.repeat(np.expand_dims(x, axis=loc), n, axis=loc)


def repeated_box(box_space, n, loc):
    return spaces.Box(
        low=stack_repeated(box_space.low, n, loc),
        high=stack_repeated(box_space.high, n, loc),
        shape=box_space.shape[:loc] + (n,) + box_space.shape[loc:],
        dtype=box_space.dtype,
    )


def repeated_space(space, n, loc=0):
    if isinstance(space, spaces.Box):
        return repeated_box(space, n, loc)
    elif isinstance(space, spaces.Dict):
        result_space = spaces.Dict()
        for key, value in space.items():
            result_space[key] = repeated_space(value, n, loc)
        return result_space
    elif isinstance(space, spaces.Discrete):
        return spaces.MultiDiscrete([[space.n] for _ in range(n)])
    elif isinstance(space, spaces.Text):  # For language, we don't repeat and only keep the last one
        return space
    else:
        raise RuntimeError(f"Unsupported space type {type(space)}")


def take_last_n(x, n):
    x = list(x)
    n = min(len(x), n)
    return np.array(x[-n:])


def dict_take_last_n(x, n):
    result = dict()
    for key, value in x.items():
        result[key] = take_last_n(value, n)
    return result


def compress_dict_list(ds, recursive=False):
    """
    Args:
        ds: list of dicts with the same keys and the same value type
        recursive: whether to recursively compress nested dictionaries
    Returns:
        dict of lists with the same keys as the dicts in ds
    """
    if not ds:
        return {}

    # Assert that ds is a list of dictionaries
    if not isinstance(ds, list):
        raise TypeError(f"Expected a list of dictionaries, but got {type(ds)}")

    if not all(isinstance(d, dict) for d in ds):
        non_dict_indices = [i for i, d in enumerate(ds) if not isinstance(d, dict)]
        raise TypeError(
            f"All elements must be dictionaries. Found non-dictionary elements at indices: {non_dict_indices}"
        )

    # Check that all dictionaries have the same keys
    keys = set(ds[0].keys())
    for i, d in enumerate(ds[1:], 1):
        if set(d.keys()) != keys:
            missing_keys = keys - set(d.keys())
            extra_keys = set(d.keys()) - keys
            error_msg = f"Dictionary at index {i} has different keys than the first dictionary."
            if missing_keys:
                error_msg += f" Missing keys: {missing_keys}."
            if extra_keys:
                error_msg += f" Extra keys: {extra_keys}."
            raise ValueError(error_msg)

    result = defaultdict(list)
    for d in ds:
        for key, value in d.items():
            result[key].append(value)

    # Convert lists to numpy arrays or recursively compress nested dictionaries
    for key, value_list in result.items():
        # Check if all values are dictionaries and recursion is enabled
        if recursive and all(isinstance(v, dict) for v in value_list):
            result[key] = compress_dict_list(value_list, recursive=True)
        else:
            try:
                result[key] = np.array(value_list)
            except Exception as e:
                raise ValueError(
                    f"Failed to convert values for key '{key}' to numpy array: {str(e)}"
                )

    return result


def aggregate(data, method: AggregateMethod = AggregateMethod.MAX):
    try:
        method = AggregateMethod(method)
    except ValueError:
        raise ValueError(
            f"Unsupported aggregate method {method!r}; "
            f"expected one of {[m.value for m in AggregateMethod]}."
        )
    if method is AggregateMethod.MAX:
        # equivalent to any
        return np.max(data)
    elif method is AggregateMethod.MIN:
        # equivalent to all
        return np.min(data)
    elif method is AggregateMethod.MEAN:
        return np.mean(data)
    elif method is AggregateMethod.SUM:
        return np.sum(data)


class MultiStepWrapper(gym.Wrapper):
    def __init__(
        self,
        env,
        contract,
        max_episode_steps=None,
        reward_agg_method: AggregateMethod = AggregateMethod.MAX,
        terminate_on_success=False,
    ):
        """
        contract: policy-resolved
          :class:`gr00t.eval._horizon_contract.PolicyHorizonSpec` carrying
          ``n_action_steps`` and the video / state ``delta_indices``.
          ``state_delta_indices`` is ``None`` for a vision-only policy. See
          ``assert_delta_indices`` for the per-array requirements.
        """
        try:
            reward_agg_method = AggregateMethod(reward_agg_method)
        except ValueError:
            raise ValueError(
                f"Unsupported reward_agg_method {reward_agg_method!r}; "
                f"expected one of {[m.value for m in AggregateMethod]}."
            )
        super().__init__(env)
        self.contract = contract
        video_delta_indices = contract.video_delta_indices_array
        state_delta_indices = contract.state_delta_indices_array
        n_action_steps = contract.n_action_steps

        # Assign action space
        self._action_space = repeated_space(env.action_space, n_action_steps)

        # Assign delta indices and horizons
        self.video_delta_indices = video_delta_indices
        self.video_horizon = len(video_delta_indices)
        self.assert_delta_indices(self.video_delta_indices, self.video_horizon)
        if state_delta_indices is not None:
            self.state_delta_indices = state_delta_indices
            self.state_horizon = len(state_delta_indices)
            self.assert_delta_indices(self.state_delta_indices, self.state_horizon)
        else:
            self.state_horizon = None
            self.state_delta_indices = None

        # Assign observation space
        self._observation_space = self.convert_observation_space(
            self.observation_space,
            self.video_horizon,
            self.state_horizon,
        )

        # Assign other attributes
        self.max_episode_steps = max_episode_steps
        self.n_action_steps = n_action_steps
        self.reward_agg_method = reward_agg_method
        self.max_steps_needed = self.get_max_steps_needed()

        self.obs = deque(maxlen=self.max_steps_needed + 1)
        self.reward = list()
        self.done = list()
        self.info = defaultdict(lambda: deque(maxlen=self.n_action_steps + 1))
        self.terminate_on_success = terminate_on_success

    def convert_observation_space(self, observation_space, video_horizon, state_horizon):
        """
        For video, the observation space will be (video_horizon,) + original shape
        For state (if not None), the observation space will be (state_horizon,) + original shape
        """
        new_observation_space = {}
        for k in observation_space.keys():
            if k.startswith("video"):
                box = observation_space[k]
                horizon = video_horizon
                new_observation_space[k] = repeated_space(box, horizon)
            elif k.startswith("state"):
                box = observation_space[k]
                if state_horizon is not None:
                    horizon = state_horizon
                else:
                    # Don't include the state in the observation space
                    continue
                new_observation_space[k] = repeated_space(box, horizon)
            elif k.startswith("annotation"):
                text = observation_space[k]
                new_observation_space[k] = text
            else:
                warnings.warn(f"Key without a prefix: {k}")
                box = observation_space[k]
                horizon = state_horizon
                new_observation_space[k] = repeated_space(box, horizon)

        return spaces.Dict(new_observation_space)

    def get_max_steps_needed(self):
        """
        Get the maximum number of steps that we need to cache.
        """
        video_max_steps_needed = (
            np.max(self.video_delta_indices) - np.min(self.video_delta_indices) + 1
        )
        if self.state_delta_indices is not None:
            state_max_steps_needed = (
                np.max(self.state_delta_indices) - np.min(self.state_delta_indices) + 1
            )
        else:
            state_max_steps_needed = 0
        return int(max(video_max_steps_needed, state_max_steps_needed))

    def assert_delta_indices(self, delta_indices: np.ndarray, horizon: int):
        # Check the length
        # (In this wrapper, this seems redundant because we get the horizon from the delta indices. But in the policy, the horizon is not derived from the delta indices but we need to make it consistent. To make the function consistent, we keep the check here.)
        assert len(delta_indices) == horizon, f"{delta_indices=}, {horizon=}"
        # All delta indices should be non-positive because there's no way to get the future observations
        assert np.all(delta_indices <= 0), f"{delta_indices=}"
        # The last delta index should be 0 because it doesn't make sense to not use the latest observation
        assert delta_indices[-1] == 0, f"{delta_indices=}"
        if len(delta_indices) > 1:
            # The step is consistent (because in real robot experiments, we actually use the dt to get the observations, which requires the step to be consistent)
            assert np.all(np.diff(delta_indices) == delta_indices[1] - delta_indices[0]), (
                f"{delta_indices=}"
            )
            # And the step is positive
            assert (delta_indices[1] - delta_indices[0]) > 0, f"{delta_indices=}"

    def reset(self, seed=None, options=None):
        """Resets the environment using kwargs."""
        obs, info = super().reset(seed=seed, options=options)

        self.obs = deque([obs] * (self.max_steps_needed + 1), maxlen=self.max_steps_needed + 1)
        self.reward = list()
        self.done = list()
        self.info = defaultdict(lambda: deque(maxlen=self.n_action_steps + 1))

        obs = self._get_obs(self.video_delta_indices, self.state_delta_indices)
        info = {k: [v] for k, v in info.items()}
        if "intermediate_signals" in info:
            # "intermediate_signals" contain the metrics for 5DC tasks to indicate language following
            info["intermediate_signals"] = {}
        return obs, info

    def step(self, action):
        """
        action: dict: key-value pairs where the values are of shape (n_action_steps,) + action_shape
        """
        states = []
        rewards = []
        dones = []
        n_env_steps = 0
        for step in range(self.n_action_steps):
            act = {}
            for key, value in action.items():
                act[key] = value[step, :]
            if len(self.done) > 0 and self.done[-1]:
                # termination
                break
            observation, reward, done, truncated, info = super().step(act)
            n_env_steps += 1
            # TODO: assign meaningful values
            env_state = {"states": [], "model": []}
            states.append(env_state["states"])
            rewards.append(reward)
            dones.append(done)
            self.obs.append(observation)
            self.reward.append(reward)
            if (self.max_episode_steps is not None) and (
                len(self.reward) >= self.max_episode_steps
            ):
                # truncation
                done = True
            self.done.append(done)
            self._add_info(info)

        observation = self._get_obs(self.video_delta_indices, self.state_delta_indices)
        reward = aggregate(self.reward, self.reward_agg_method)
        done = aggregate(self.done, AggregateMethod.MAX)
        info = dict_take_last_n(self.info, self.n_action_steps)
        states = np.array(states)
        rewards = np.array(rewards)
        dones = np.array(dones)
        info["states"] = states
        info["rewards"] = rewards
        info["model"] = env_state["model"]
        info["actions"] = action
        info["dones"] = dones
        # Inner env-steps taken this macro-step (< n_action_steps if `done` fires
        # mid-chunk); run_rollout_gymnasium_policy sums it for episode_length.
        info["n_env_steps"] = n_env_steps
        if "intermediate_signals" in info:
            # "intermediate_signals" contain the metrics for 5DC tasks to indicate language following
            # Here we turn a list of dicts into a dict of lists
            """
            Example of the ultimate format of `info["intermediate_signals"]`:
            {
                'grasp_obj': [True, ..., True],
                'grasp_distractor_obj': [False, ..., False],
                'gripper_obj_dist': [0.0004638251298563212, ..., 0.0004638251298563212],
                'gripper_distractor_dist': [0.0023107511879928433, ..., 0.0023107511879928433]
            }
            The length is `n_action_steps`.
            """
            info["intermediate_signals"] = compress_dict_list(list(info["intermediate_signals"]))

        if self.terminate_on_success and any(info["success"]):
            # Terminate after this step.
            done = True

        return observation, reward, done, truncated, info

    def _get_obs(self, video_delta_indices, state_delta_indices):
        """
        Output:
        For video: (video_horizon,) + obs_shape
        For state (if not None): (state_horizon,) + obs_shape
        """
        assert len(self.obs) > 0
        if isinstance(self.observation_space, spaces.Dict):
            result = dict()
            for key in self.observation_space.keys():
                if key.startswith("video"):
                    """
                    NOTE:
                      We need to subtract 1 because video_delta_indices is 0-indexed.
                      E.g., video_delta_indices = np.array([-4, -3, -2, -1, 0])
                      Then when we select the observation,
                        it should be [obs[-5], obs[-4], obs[-3], obs[-2], obs[-1]]
                      (i.e., the latest observation is at the last index)
                    """
                    delta_indices = video_delta_indices - 1
                    this_obs = [self.obs[i][key] for i in delta_indices]
                    result[key] = np.stack(this_obs, axis=0)
                elif key.startswith("state"):
                    if state_delta_indices is not None:
                        delta_indices = state_delta_indices - 1
                    else:
                        raise ValueError(
                            f"state_delta_indices is None but `state` is still in the {self.observation_space=}"
                        )
                    this_obs = [self.obs[i][key] for i in delta_indices]
                    result[key] = np.stack(this_obs, axis=0)
                elif key.startswith("annotation"):
                    result[key] = self.obs[-1][key]
                else:
                    if state_delta_indices is not None:
                        delta_indices = state_delta_indices - 1
                    else:
                        raise ValueError(
                            f"state_delta_indices is None but `state` is still in the {self.observation_space=}"
                        )
                    this_obs = [self.obs[i][key] for i in delta_indices]
                    result[key] = np.stack(this_obs, axis=0)
            return result
        else:
            raise RuntimeError(f"Unsupported space type: {type(self.observation_space)=}")

    def _add_info(self, info):
        for key, value in info.items():
            self.info[key].append(value)

    def get_rewards(self):
        return self.reward

    def get_attr(self, name):
        return getattr(self, name)

    def get_infos(self):
        result = dict()
        for k, v in self.info.items():
            result[k] = list(v)
        return result