chore: import upstream snapshot with attribution

rllib/utils/postprocessing/episodes.py

@@ -0,0 +1,144 @@
+from typing import List, Tuple
+
+import numpy as np
+
+from ray.rllib.env.single_agent_episode import SingleAgentEpisode
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+def add_one_ts_to_episodes_and_truncate(episodes: List[SingleAgentEpisode]):
+    """Adds an artificial timestep to an episode at the end.
+
+    In detail: The last observations, infos, actions, and all `extra_model_outputs`
+    will be duplicated and appended to each episode's data. An extra 0.0 reward
+    will be appended to the episode's rewards. The episode's timestep will be
+    increased by 1. Also, adds the truncated=True flag to each episode if the
+    episode is not already done (terminated or truncated).
+
+    Useful for value function bootstrapping, where it is required to compute a
+    forward pass for the very last timestep within the episode,
+    i.e. using the following input dict: {
+      obs=[final obs],
+      state=[final state output],
+      prev. reward=[final reward],
+      etc..
+    }
+
+    Args:
+        episodes: The list of SingleAgentEpisode objects to extend by one timestep
+            and add a truncation flag if necessary.
+
+    Returns:
+        A list of the original episodes' truncated values (so the episodes can be
+        properly restored later into their original states).
+    """
+    orig_truncateds = []
+    for episode in episodes:
+        orig_truncateds.append(episode.is_truncated)
+
+        # Add timestep.
+        episode.t += 1
+        # Use the episode API that allows appending (possibly complex) structs
+        # to the data.
+        episode.observations.append(episode.observations[-1])
+        episode.infos.append(episode.infos[-1])
+        episode.actions.append(episode.actions[-1])
+        episode.rewards.append(0.0)
+        for v in episode.extra_model_outputs.values():
+            v.append(v[-1])
+        # Artificially make this episode truncated for the upcoming GAE
+        # computations.
+        if not episode.is_done:
+            episode.is_truncated = True
+        # Validate to make sure, everything is in order.
+        episode.validate()
+
+    return orig_truncateds
+
+
+@DeveloperAPI
+def remove_last_ts_from_data(
+    episode_lens: List[int],
+    *data: Tuple[np._typing.NDArray],
+) -> Tuple[np._typing.NDArray]:
+    """Removes the last timesteps from each given data item.
+
+    Each item in data is a concatenated sequence of episodes data.
+    For example if `episode_lens` is [2, 4], then data is a shape=(6,)
+    ndarray. The returned corresponding value will have shape (4,), meaning
+    both episodes have been shortened by exactly one timestep to 1 and 3.
+
+    ..testcode::
+
+        from ray.rllib.algorithms.ppo.ppo_learner import PPOLearner
+        import numpy as np
+
+        unpadded = PPOLearner._remove_last_ts_from_data(
+            [5, 3],
+            np.array([0, 1, 2, 3, 4,  0, 1, 2]),
+        )
+        assert (unpadded[0] == [0, 1, 2, 3, 0, 1]).all()
+
+        unpadded = PPOLearner._remove_last_ts_from_data(
+            [4, 2, 3],
+            np.array([0, 1, 2, 3,  0, 1,  0, 1, 2]),
+            np.array([4, 5, 6, 7,  2, 3,  3, 4, 5]),
+        )
+        assert (unpadded[0] == [0, 1, 2,  0,  0, 1]).all()
+        assert (unpadded[1] == [4, 5, 6,  2,  3, 4]).all()
+
+    Args:
+        episode_lens: A list of current episode lengths. The returned
+            data will have the same lengths minus 1 timestep.
+        data: A tuple of data items (np.ndarrays) representing concatenated episodes
+            to be shortened by one timestep per episode.
+            Note that only arrays with `shape=(n,)` are supported! The
+            returned data will have `shape=(n-len(episode_lens),)` (each
+            episode gets shortened by one timestep).
+
+    Returns:
+        A tuple of new data items shortened by one timestep.
+    """
+    # Figure out the new slices to apply to each data item based on
+    # the given episode_lens.
+    slices = []
+    sum = 0
+    for len_ in episode_lens:
+        slices.append(slice(sum, sum + len_ - 1))
+        sum += len_
+    # Compiling return data by slicing off one timestep at the end of
+    # each episode.
+    ret = []
+    for d in data:
+        ret.append(np.concatenate([d[s] for s in slices]))
+    return tuple(ret) if len(ret) > 1 else ret[0]
+
+
+@DeveloperAPI
+def remove_last_ts_from_episodes_and_restore_truncateds(
+    episodes: List[SingleAgentEpisode],
+    orig_truncateds: List[bool],
+) -> None:
+    """Reverts the effects of `_add_ts_to_episodes_and_truncate`.
+
+    Args:
+        episodes: The list of SingleAgentEpisode objects to extend by one timestep
+            and add a truncation flag if necessary.
+        orig_truncateds: A list of the original episodes' truncated values to be
+            applied to the `episodes`.
+    """
+
+    # Fix all episodes.
+    for episode, orig_truncated in zip(episodes, orig_truncateds):
+        # Reduce timesteps by 1.
+        episode.t -= 1
+        # Remove all extra timestep data from the episode's buffers.
+        episode.observations.pop()
+        episode.infos.pop()
+        episode.actions.pop()
+        episode.rewards.pop()
+        for v in episode.extra_model_outputs.values():
+            v.pop()
+        # Fix the truncateds flag again.
+        episode.is_truncated = orig_truncated

rllib/utils/postprocessing/tests/test_value_predictions.py

@@ -0,0 +1,48 @@
+import unittest
+
+from ray.rllib.utils.postprocessing.value_predictions import extract_bootstrapped_values
+from ray.rllib.utils.test_utils import check
+
+
+class TestPostprocessing(unittest.TestCase):
+    def test_extract_bootstrapped_values(self):
+        """Tests, whether the extract_bootstrapped_values utility works properly."""
+
+        # Fake vf_preds sequence.
+        # Spaces = denote (elongated-by-one-artificial-ts) episode boundaries.
+        # digits = timesteps within the actual episode.
+        # [lower case letters] = bootstrap values at episode truncations.
+        # '-' = bootstrap values at episode terminals (these values are simply zero).
+        sequence = "012345678a 01234A 0- 0123456b 01c 012- 012345e 012-"
+        sequence = sequence.replace(" ", "")
+        sequence = list(sequence)
+        # The actual, non-elongated, episode lengths.
+        episode_lengths = [9, 5, 1, 7, 2, 3, 6, 3]
+        T = 4
+        result = extract_bootstrapped_values(
+            vf_preds=sequence,
+            episode_lengths=episode_lengths,
+            T=T,
+        )
+        check(result, [4, 8, 3, 1, 5, "c", 1, 5, "-"])
+
+        # Another example.
+        sequence = "0123a 012345b 01234567- 012- 012- 012- 012345- 0123456c"
+        sequence = sequence.replace(" ", "")
+        sequence = list(sequence)
+        episode_lengths = [4, 6, 8, 3, 3, 3, 6, 7]
+        T = 5
+        result = extract_bootstrapped_values(
+            vf_preds=sequence,
+            episode_lengths=episode_lengths,
+            T=T,
+        )
+        check(result, [1, "b", 5, 2, 1, 3, 2, "c"])
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))

rllib/utils/postprocessing/value_predictions.py

@@ -0,0 +1,125 @@
+import numpy as np
+
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+def compute_value_targets(
+    values,
+    rewards,
+    terminateds,
+    truncateds,
+    gamma: float,
+    lambda_: float,
+):
+    """Computes GAE value targets given vf predictions and rewards.
+
+    Convention (Gymnasium-aligned, matches ``AddOneTsToEpisodesAndTruncate``):
+      ``terminateds[t] = True``  =>  no s_{t+1}; gate t -> t+1 bootstrap.
+      ``truncateds[t]  = True``  =>  step t ends an episode chunk; V(s_{t+1})
+                                     remains a valid bootstrap, but GAE must
+                                     not propagate across the boundary.
+
+    Advantages = targets - vf_predictions.
+
+    See https://pseudo-rnd-thoughts.github.io/blog/visualising-gae/ for visualisation.
+    """
+    # 1 if the transition t -> t+1 exists (not a terminal at t), else 0.
+    non_terminal = 1.0 - terminateds
+    # 1 if GAE may propagate from t+1 back into t, else 0. Both terminal and
+    # chunk-boundary steps stop the recursion.
+    propagate = non_terminal * (1.0 - truncateds)
+
+    # V(s_{t+1}) per timestep. The trailing 0.0 is a dummy: the corresponding
+    # td_residual is masked out downstream by `loss_mask`, and the recursion
+    # carrying it is gated by `propagate`.
+    next_state_values = np.append(values[1:], 0.0)
+
+    # TD residual: delta_t = r_t + gamma * (1 - terminated_t) * V(s_{t+1}) - V(s_t)
+    # Truncation does NOT zero the bootstrap -- V(s_{t+1}) is a valid
+    # prediction at a truncation boundary.
+    td_residuals = rewards + gamma * non_terminal * next_state_values - values
+
+    # GAE backward recursion. `running_advantage` carries advantage[t+1] into
+    # iteration t and is killed at terminal / truncation boundaries by
+    # `propagate`.
+    advantages = np.zeros_like(rewards, dtype=np.float32)
+    running_advantage = 0.0
+    for t in reversed(range(td_residuals.shape[0])):
+        running_advantage = (
+            td_residuals[t] + gamma * lambda_ * propagate[t] * running_advantage
+        )
+        advantages[t] = running_advantage
+
+    # target_t = advantage_t + V(s_t).
+    return (advantages + values).astype(np.float32)
+
+
+def extract_bootstrapped_values(vf_preds, episode_lengths, T):
+    """Returns a bootstrapped value batch given value predictions.
+
+    Note that the incoming value predictions must have happened over (artificially)
+    elongated episodes (by 1 timestep at the end). This way, we can either extract the
+    `vf_preds` at these extra timesteps (as "bootstrap values") or skip over them
+    entirely if they lie in the middle of the T-slices.
+
+    For example, given an episodes structure like this:
+    01234a 0123456b 01c 012- 0123e 012-
+    where each episode is separated by a space and goes from 0 to n and ends in an
+    artificially elongated timestep (denoted by 'a', 'b', 'c', '-', or 'e'), where '-'
+    means that the episode was terminated and the bootstrap value at the end should be
+    zero and 'a', 'b', 'c', etc.. represent truncated episode ends with computed vf
+    estimates.
+    The output for the above sequence (and T=4) should then be:
+    4 3 b 2 3 -
+
+    Args:
+        vf_preds: The computed value function predictions over the artificially
+            elongated episodes (by one timestep at the end).
+        episode_lengths: The original (correct) episode lengths, NOT counting the
+            artificially added timestep at the end.
+        T: The size of the time dimension by which to slice the data. Note that the
+            sum of all episode lengths (`sum(episode_lengths)`) must be dividable by T.
+
+    Returns:
+        The batch of bootstrapped values.
+    """
+    bootstrapped_values = []
+    if sum(episode_lengths) % T != 0:
+        raise ValueError(
+            "Can only extract bootstrapped values if the sum of episode lengths "
+            f"({sum(episode_lengths)}) is dividable by the given T ({T})!"
+        )
+
+    # Loop over all episode lengths and collect bootstrap values.
+    # Do not alter incoming `episode_lengths` list.
+    episode_lengths = episode_lengths[:]
+    i = -1
+    while i < len(episode_lengths) - 1:
+        i += 1
+        eps_len = episode_lengths[i]
+        # We can make another T-stride inside this episode ->
+        # - Use a vf prediction within the episode as bootstrapped value.
+        # - "Fix" the episode_lengths array and continue within the same episode.
+        if T < eps_len:
+            bootstrapped_values.append(vf_preds[T])
+            vf_preds = vf_preds[T:]
+            episode_lengths[i] -= T
+            i -= 1
+        # We can make another T-stride inside this episode, but will then be at the end
+        # of it ->
+        # - Use the value function prediction at the artificially added timestep
+        #   as bootstrapped value.
+        # - Skip the additional timestep at the end and ,ove on with next episode.
+        elif T == eps_len:
+            bootstrapped_values.append(vf_preds[T])
+            vf_preds = vf_preds[T + 1 :]
+        # The episode fits entirely into the T-stride ->
+        # - Move on to next episode ("fix" its length by make it seemingly longer).
+        else:
+            # Skip bootstrap value of current episode (not needed).
+            vf_preds = vf_preds[1:]
+            # Make next episode seem longer.
+            episode_lengths[i + 1] += eps_len
+
+    return np.array(bootstrapped_values)

rllib/utils/postprocessing/zero_padding.py

@@ -0,0 +1,270 @@
+from collections import deque
+from typing import List, Tuple, Union
+
+import numpy as np
+import tree  # pip install dm_tree
+
+from ray.rllib.utils.spaces.space_utils import BatchedNdArray, batch
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+def create_mask_and_seq_lens(episode_len: int, T: int) -> Tuple[List, List]:
+    """Creates loss mask and a seq_lens array, given an episode length and T.
+
+    Args:
+        episode_lens: A list of episode lengths to infer the loss mask and seq_lens
+            array from.
+        T: The maximum number of timesteps in each "row", also known as the maximum
+            sequence length (max_seq_len). Episodes are split into chunks that are at
+            most `T` long and remaining timesteps will be zero-padded (and masked out).
+
+    Returns:
+         Tuple consisting of a) list of the loss masks to use (masking out areas that
+         are past the end of an episode (or rollout), but had to be zero-added due to
+         the added extra time rank (of length T) and b) the list of sequence lengths
+         resulting from splitting the given episodes into chunks of at most `T`
+         timesteps.
+    """
+    mask = []
+    seq_lens = []
+
+    len_ = min(episode_len, T)
+    seq_lens.append(len_)
+    row = np.array([1] * len_ + [0] * (T - len_), np.bool_)
+    mask.append(row)
+
+    # Handle sequence lengths greater than T.
+    overflow = episode_len - T
+    while overflow > 0:
+        len_ = min(overflow, T)
+        seq_lens.append(len_)
+        extra_row = np.array([1] * len_ + [0] * (T - len_), np.bool_)
+        mask.append(extra_row)
+        overflow -= T
+
+    return mask, seq_lens
+
+
+@DeveloperAPI
+def split_and_zero_pad(
+    item_list: List[Union[BatchedNdArray, np._typing.NDArray, float]],
+    max_seq_len: int,
+) -> List[np._typing.NDArray]:
+    """Splits the contents of `item_list` into a new list of ndarrays and returns it.
+
+    In the returned list, each item is one ndarray of len (axis=0) `max_seq_len`.
+    The last item in the returned list may be (right) zero-padded, if necessary, to
+    reach `max_seq_len`.
+
+    If `item_list` contains one or more `BatchedNdArray` (instead of individual
+    items), these will be split accordingly along their axis=0 to yield the returned
+    structure described above.
+
+    .. testcode::
+
+        from ray.rllib.utils.postprocessing.zero_padding import (
+            BatchedNdArray,
+            split_and_zero_pad,
+        )
+        from ray.rllib.utils.test_utils import check
+
+        # Simple case: `item_list` contains individual floats.
+        check(
+            split_and_zero_pad([0, 1, 2, 3, 4, 5, 6, 7], 5),
+            [[0, 1, 2, 3, 4], [5, 6, 7, 0, 0]],
+        )
+
+        # `item_list` contains BatchedNdArray (ndarrays that explicitly declare they
+        # have a batch axis=0).
+        check(
+            split_and_zero_pad([
+                BatchedNdArray([0, 1]),
+                BatchedNdArray([2, 3, 4, 5]),
+                BatchedNdArray([6, 7, 8]),
+            ], 5),
+            [[0, 1, 2, 3, 4], [5, 6, 7, 8, 0]],
+        )
+
+    Args:
+        item_list: A list of individual items or BatchedNdArrays to be split into
+            `max_seq_len` long pieces (the last of which may be zero-padded).
+        max_seq_len: The maximum length of each item in the returned list.
+
+    Returns:
+        A list of np.ndarrays (all of length `max_seq_len`), which contains the same
+        data as `item_list`, but split into sub-chunks of size `max_seq_len`.
+        The last item in the returned list may be zero-padded, if necessary.
+    """
+    zero_element = tree.map_structure(
+        lambda s: np.zeros_like([s[0]] if isinstance(s, BatchedNdArray) else s),
+        item_list[0],
+    )
+
+    # The replacement list (to be returned) for `items_list`.
+    # Items list contains n individual items.
+    # -> ret will contain m batched rows, where m == n // T and the last row
+    # may be zero padded (until T).
+    ret = []
+
+    # List of the T-axis item, collected to form the next row.
+    current_time_row = []
+    current_t = 0
+
+    item_list = deque(item_list)
+    while len(item_list) > 0:
+        item = item_list.popleft()
+        t = max_seq_len - current_t
+
+        # In case `item` is a complex struct.
+        item_flat = tree.flatten(item)
+        item_list_append = []
+        current_time_row_flat_items = []
+        add_to_current_t = 0
+
+        for itm in item_flat:
+            # `itm` is already a batched np.array: Split if necessary.
+            if isinstance(itm, BatchedNdArray):
+                current_time_row_flat_items.append(itm[:t])
+                if len(itm) <= t:
+                    add_to_current_t = len(itm)
+                else:
+                    add_to_current_t = t
+                    item_list_append.append(itm[t:])
+            # `itm` is a single item (no batch axis): Append and continue with next
+            # item.
+            else:
+                current_time_row_flat_items.append(itm)
+                add_to_current_t = 1
+
+        current_t += add_to_current_t
+        current_time_row.append(tree.unflatten_as(item, current_time_row_flat_items))
+        if item_list_append:
+            item_list.appendleft(tree.unflatten_as(item, item_list_append))
+
+        # `current_time_row` is "full" (max_seq_len): Append as ndarray (with batch
+        # axis) to `ret`.
+        if current_t == max_seq_len:
+            ret.append(
+                batch(
+                    current_time_row,
+                    individual_items_already_have_batch_dim="auto",
+                )
+            )
+            current_time_row = []
+            current_t = 0
+
+    # `current_time_row` is unfinished: Pad, if necessary and append to `ret`.
+    if current_t > 0 and current_t < max_seq_len:
+        current_time_row.extend([zero_element] * (max_seq_len - current_t))
+        ret.append(
+            batch(current_time_row, individual_items_already_have_batch_dim="auto")
+        )
+
+    return ret
+
+
+@DeveloperAPI
+def split_and_zero_pad_n_episodes(
+    nd_array: np._typing.NDArray,
+    episode_lens: List[int],
+    max_seq_len: int,
+) -> List[np._typing.NDArray]:
+    """Splits and zero-pads a single np.ndarray based on episode lens and a maxlen.
+
+    Args:
+        nd_array: The single np.ndarray to be split into n chunks, based on the given
+            `episode_lens` and the `max_seq_len` argument. For example, if `nd_array`
+            has a batch dimension (axis 0) of 21, `episode_lens` is [15, 3, 3], and
+            `max_seq_len` is 6, then the returned list would have np.ndarrays in it of
+            batch dimensions (axis 0): [6, 6, 6 (zero-padded), 6 (zero-padded),
+            6 (zero-padded)].
+            Note that this function doesn't work on nested data, such as dicts of
+            ndarrays.
+        episode_lens: A list of episode lengths along which to split and zero-pad the
+            given `nd_array`.
+        max_seq_len: The maximum sequence length to split at (and zero-pad).
+
+    Returns: A list of n np.ndarrays, resulting from splitting and zero-padding the
+        given `nd_array`.
+    """
+    ret = []
+
+    cursor = 0
+    for episode_len in episode_lens:
+        items = BatchedNdArray(nd_array[cursor : cursor + episode_len])
+        ret.extend(split_and_zero_pad([items], max_seq_len))
+        cursor += episode_len
+
+    return ret
+
+
+@DeveloperAPI
+def unpad_data_if_necessary(
+    episode_lens: List[int],
+    data: np._typing.NDArray,
+) -> np._typing.NDArray:
+    """Removes right-side zero-padding from data based on `episode_lens`.
+
+    ..testcode::
+
+        from ray.rllib.utils.postprocessing.zero_padding import unpad_data_if_necessary
+        import numpy as np
+
+        unpadded = unpad_data_if_necessary(
+            episode_lens=[4, 2],
+            data=np.array([
+                [2, 4, 5, 3, 0, 0, 0, 0],
+                [-1, 3, 0, 0, 0, 0, 0, 0],
+            ]),
+        )
+        assert (unpadded == [2, 4, 5, 3, -1, 3]).all()
+
+        unpadded = unpad_data_if_necessary(
+            episode_lens=[1, 5],
+            data=np.array([
+                [2, 0, 0, 0, 0],
+                [-1, -2, -3, -4, -5],
+            ]),
+        )
+        assert (unpadded == [2, -1, -2, -3, -4, -5]).all()
+
+    Args:
+        episode_lens: A list of actual episode lengths.
+        data: A 2D np.ndarray with right-side zero-padded rows.
+
+    Returns:
+        A 1D np.ndarray resulting from concatenation of the un-padded
+        input data along the 0-axis.
+    """
+    # If data des NOT have time dimension, return right away.
+    if len(data.shape) == 1:
+        return data
+
+    # Assert we only have B and T dimensions (meaning this function only operates
+    # on single-float data, such as value function predictions, advantages, or rewards).
+    assert len(data.shape) == 2
+
+    new_data = []
+    row_idx = 0
+
+    T = data.shape[1]
+    for len_ in episode_lens:
+        # Calculate how many full rows this array occupies and how many elements are
+        # in the last, potentially partial row.
+        num_rows, col_idx = divmod(len_, T)
+
+        # If the array spans multiple full rows, fully include these rows.
+        for i in range(num_rows):
+            new_data.append(data[row_idx])
+            row_idx += 1
+
+        # If there are elements in the last, potentially partial row, add this
+        # partial row as well.
+        if col_idx > 0:
+            new_data.append(data[row_idx, :col_idx])
+
+            # Move to the next row for the next array (skip the zero-padding zone).
+            row_idx += 1
+
+    return np.concatenate(new_data)