File size: 6,637 Bytes
fbd9366 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 | """
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)
|