def broadcast_to_axes(axes: Union[AxesLike, None], *arrays: AxesParams):
"""Broadcast ``arrays`` to the length of ``axes``. Axes are inferred from the arrays if necessary."""
if not arrays:
raise ValueError('No arrays provided.')
arrays = lmap(np.atleast_1d, arrays)
lengths = lmap(len, arrays)
if axes is None:
axes = list(range(-max(lengths), 0))
axes = check_axes(axes)
if not all(len(axes) == x or x == 1 for x in lengths):
raise ValueError(
f'Axes and arrays are not broadcastable: {len(axes)} vs {join(lengths)}.'
)
arrays = [np.repeat(x, len(axes) // len(x), 0) for x in arrays]
return (axes, *arrays)
def __init__(self,
layer: Callable,
downsample: Union[nn.Module, Callable],
upsample: Union[nn.Module, Callable],
merge: Callable,
structure: Sequence[Sequence[Union[Sequence[int],
nn.Module]]],
last_level: bool = True,
**kwargs):
super().__init__()
def build_level(path):
if isinstance(path, nn.Module):
return path
elif not isinstance(path, Sequence) or not all(
isinstance(x, int) for x in path):
raise ValueError('The passed `structure` is not valid.')
return ConsistentSequential(layer, path, **kwargs)
def make_up_down(o):
if not isinstance(o, nn.Module):
o = o()
return o
*levels, bridge = structure
# handling the case [[...]]
if len(bridge) == 1 and isinstance(bridge[0], Sequence):
bridge = bridge[0]
self.bridge = build_level(bridge)
self.merge = merge
self.last_level = last_level
self.downsample = nn.ModuleList(
[make_up_down(downsample) for _ in levels])
self.upsample = nn.ModuleList([make_up_down(upsample) for _ in levels])
# group branches
branches = []
for paths in zip_equal(*structure[:-1]):
branches.append(nn.ModuleList(lmap(build_level, paths)))
if len(branches) not in [2, 3]:
raise ValueError(
f'Expected 2 or 3 branches, but {len(branches)} provided.')
self.down_path, self.up_path = branches[0], branches[-1]
# add middle branch if needed
if len(branches) == 2:
self.middle_path = [identity] * len(self.down_path)
else:
self.middle_path = branches[1]
def describe_dqn(memory: Memory, agent, gamma: float = 1):
state = memory.sample_episode().state(0)
rewards = [e.rewards() for e in memory.episodes()]
sizes = np.array(lmap(len, rewards))
discounted = np.array([discount_rewards(r, gamma) for r in rewards])
summed = np.array([discount_rewards(r, 1) for r in rewards])
q_values = get_q_values(state, agent)
return {
'mean reward': discounted.mean(), 'rewards': discounted[None],
'mean reward sum': summed.mean(), 'rewards sum': summed[None],
'mean size': sizes.mean(), 'sizes': sizes[None],
'q_max': q_values.max(), 'q_min': q_values.min(),
}
def wrapper(*arrays):
return np.stack(
lmap(unpack_args(predict), iterate_slices(*arrays, axis=-1)), -1)
def uniform(shape):
return np.array(lmap(np.random.randint, np.atleast_1d(shape)))