def forward(self, inputs, state=()):
"""
Args:
inputs (nested Tensor):
"""
# call super to preprocess inputs
z, state = super().forward(inputs, state)
if self._img_encoding_net is not None:
z, _ = self._img_encoding_net(z)
if alf.summary.should_summarize_output():
name = ('summarize_output/' + self.name + '.fc.0.' +
'input_norm.' + common.exe_mode_name())
alf.summary.scalar(name=name,
data=torch.mean(
z.norm(dim=list(range(1, z.ndim)))))
i = 0
for fc in self._fc_layers:
z = fc(z)
if alf.summary.should_summarize_output():
name = ('summarize_output/' + self.name + '.fc.' + str(i) +
'.output_norm.' + common.exe_mode_name())
alf.summary.scalar(name=name,
data=torch.mean(
z.norm(dim=list(range(1, z.ndim)))))
i += 1
if self._reshape_output:
z = z.reshape(z.shape[0], *self._output_spec.shape)
return z, state
def forward(self, observation, state=()):
"""Computes action given an observation.
Args:
inputs: A tensor consistent with ``input_tensor_spec``
state: empty for API consistent with ``ActorRNNNetwork``
Returns:
tuple:
- action (torch.Tensor): a tensor consistent with ``action_spec``
- state: empty
"""
observation, state = super().forward(observation, state)
encoded_obs, _ = self._encoding_net(observation)
actions = []
i = 0
for layer, spec in zip(self._action_layers, self._flat_action_spec):
pre_activation = layer(encoded_obs)
action = self._squashing_func(pre_activation)
action = spec_utils.scale_to_spec(action, spec)
if alf.summary.should_summarize_output():
alf.summary.scalar(
name='summarize_output/' + self.name + '.action_layer.' +
str(i) + '.pre_activation.output_norm.' +
common.exe_mode_name(),
data=torch.mean(
pre_activation.norm(
dim=list(range(1, pre_activation.ndim)))))
a_name = ('summarize_output/' + self.name + '.action_layer.' +
str(i) + '.action.output_norm.' +
common.exe_mode_name())
alf.summary.scalar(
name=a_name,
data=torch.mean(
action.norm(dim=list(range(1, action.ndim)))))
actions.append(action)
i += 1
output_actions = nest.pack_sequence_as(self._action_spec, actions)
return output_actions, state
def update(self, tensor):
"""Update the statistics given a new tensor.
"""
if self._mean_averager:
self._mean_averager.update(tensor)
sqr_tensor = alf.nest.map_structure(math_ops.square, tensor)
self._m2_averager.update(sqr_tensor)
if alf.summary.should_record_summaries():
suffix = common.exe_mode_name()
def _reduce_along_batch_dims(x, mean, op):
spec = TensorSpec.from_tensor(mean)
bs = alf.layers.BatchSquash(get_outer_rank(x, spec))
x = bs.flatten(x)
x = op(x, dim=0)[0]
return x
def _summary(name, val):
with alf.summary.scope(self._name):
if val.ndim == 0:
alf.summary.scalar(name + "." + suffix, val)
elif (val.shape[0] < self.MAX_DIMS_TO_OUTPUT
and alf.summary.should_summarize_output()):
for i in range(val.shape[0]):
alf.summary.scalar(
name + "_" + str(i) + "." + suffix, val[i])
else:
alf.summary.scalar(name + ".min." + suffix, val.min())
alf.summary.scalar(name + ".max." + suffix, val.max())
def _summarize_all(path, t, m2, m=None):
if path:
path += "."
_summary(path + "tensor.batch_min",
_reduce_along_batch_dims(t, m2, torch.min))
_summary(path + "tensor.batch_max",
_reduce_along_batch_dims(t, m2, torch.max))
if m is not None:
_summary(path + "mean", m)
_summary(path + "var", m2 - math_ops.square(m))
else:
_summary(path + "second_moment", m2)
if self._mean_averager:
alf.nest.py_map_structure_with_path(_summarize_all, tensor,
self._m2_averager.get(),
self._mean_averager.get())
else:
alf.nest.py_map_structure_with_path(_summarize_all, tensor,
self._m2_averager.get())