groot/vla/common/utils/data_structure/tree_utils.py

"""
Utils to handle nested data structures

Install dm_tree first:
https://tree.readthedocs.io/en/latest/api.html
"""

import collections
from typing import Any, Iterable, List, Tuple, TypeVar

import numpy as np

try:
    import tree

except ImportError:
    raise ImportError("Please install dm_tree first: `pip install dm_tree`")


def is_sequence(obj):
    """
    Returns:
      True if the sequence is a collections.Sequence and not a string.
    """
    return isinstance(obj, collections.abc.Sequence) and not isinstance(obj, str)


def is_mapping(obj):
    """
    Returns:
      True if the sequence is a collections.Mapping
    """
    return isinstance(obj, collections.abc.Mapping)


def tree_value_at_path(obj, paths: Tuple):
    try:
        for p in paths:
            obj = obj[p]
        return obj
    except Exception as e:
        raise ValueError(f"{e}\n\n-- Incorrect nested path {paths} for object: {obj}.")


def tree_assign_at_path(obj, paths: Tuple, value):
    try:
        for p in paths[:-1]:
            obj = obj[p]
        if len(paths) > 0:
            obj[paths[-1]] = value
    except Exception as e:
        raise ValueError(f"{e}\n\n-- Incorrect nested path {paths} for object: {obj}.")


def copy_non_leaf(obj):
    """
    Deepcopy the nested structure, but does NOT copy the leaf values like Tensors
    """
    return tree.map_structure(lambda x: x, obj)


# =======================================================================
# Copyright 2018 DeepMind Technologies Limited. All rights reserved.
#
# 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.

# Tensor framework-agnostic utilities for manipulating nested structures.

ElementType = TypeVar("ElementType")


def fast_map_structure(func, *structure):
    """Faster map_structure implementation which skips some error checking."""
    flat_structure = (tree.flatten(s) for s in structure)
    entries = zip(*flat_structure)
    # Arbitrarily choose one of the structures of the original sequence (the last)
    # to match the structure for the flattened sequence.
    return tree.unflatten_as(structure[-1], [func(*x) for x in entries])


def stack_sequence_fields(sequence: Iterable[ElementType]) -> ElementType:
    """Stacks a list of identically nested objects.

    This takes a sequence of identically nested objects and returns a single
    nested object whose ith leaf is a stacked numpy array of the corresponding
    ith leaf from each element of the sequence.

    For example, if `sequence` is:

    ```python
    [{
          'action': np.array([1.0]),
          'observation': (np.array([0.0, 1.0, 2.0]),),
          'reward': 1.0
     }, {
          'action': np.array([0.5]),
          'observation': (np.array([1.0, 2.0, 3.0]),),
          'reward': 0.0
     }, {
          'action': np.array([0.3]),1
          'observation': (np.array([2.0, 3.0, 4.0]),),
          'reward': 0.5
     }]
    ```

    Then this function will return:

    ```python
    {
        'action': np.array([....])         # array shape = [3 x 1]
        'observation': (np.array([...]),)  # array shape = [3 x 3]
        'reward': np.array([...])          # array shape = [3]
    }
    ```

    Note that the 'observation' entry in the above example has two levels of
    nesting, i.e it is a tuple of arrays.

    Args:
      sequence: a list of identically nested objects.

    Returns:
      A nested object with numpy.

    Raises:
      ValueError: If `sequence` is an empty sequence.
    """
    # Handle empty input sequences.
    if not sequence:
        raise ValueError("Input sequence must not be empty")

    # Default to asarray when arrays don't have the same shape to be compatible
    # with old behaviour.
    try:
        return fast_map_structure(lambda *values: np.stack(values), *sequence)
    except ValueError:
        return fast_map_structure(lambda *values: np.asarray(values), *sequence)


def unstack_sequence_fields(struct: ElementType, batch_size: int) -> List[ElementType]:
    """Converts a struct of batched arrays to a list of structs.

    This is effectively the inverse of `stack_sequence_fields`.

    Args:
      struct: An (arbitrarily nested) structure of arrays.
      batch_size: The length of the leading dimension of each array in the struct.
        This is assumed to be static and known.

    Returns:
      A list of structs with the same structure as `struct`, where each leaf node
       is an unbatched element of the original leaf node.
    """

    return [tree.map_structure(lambda s, i=i: s[i], struct) for i in range(batch_size)]


def broadcast_structures(*args: Any) -> Any:
    """Returns versions of the arguments that give them the same nested structure.

    Any nested items in *args must have the same structure.

    Any non-nested item will be replaced with a nested version that shares that
    structure. The leaves will all be references to the same original non-nested
    item.

    If all *args are nested, or all *args are non-nested, this function will
    return *args unchanged.

    Example:
    ```
    a = ('a', 'b')
    b = 'c'
    tree_a, tree_b = broadcast_structure(a, b)
    tree_a
    > ('a', 'b')
    tree_b
    > ('c', 'c')
    ```

    Args:
      *args: A Sequence of nested or non-nested items.

    Returns:
      `*args`, except with all items sharing the same nest structure.
    """
    if not args:
        return

    reference_tree = None
    for arg in args:
        if tree.is_nested(arg):
            reference_tree = arg
            break

    if reference_tree is None:
        reference_tree = args[0]

    def mirror_structure(value, reference_tree):
        if tree.is_nested(value):
            # Use check_types=True so that the types of the trees we construct aren't
            # dependent on our arbitrary choice of which nested arg to use as the
            # reference_tree.
            tree.assert_same_structure(value, reference_tree, check_types=True)
            return value
        else:
            return tree.map_structure(lambda _: value, reference_tree)

    return tuple(mirror_structure(arg, reference_tree) for arg in args)