tf_agents.networks.NestFlatten

Returns a Keras layer that takes a nest of inputs, and returns a list.

View aliases

Main aliases

tf_agents.networks.nest_map.NestFlatten

tf_agents.networks.NestFlatten() -> tf.keras.layers.Layer

Useful in combination with NestMap to combine processed inputs:

# Process inputs in dictionary {"inp1": ..., "inp2": ...}, then
# flatten the resulting tensors into a list, and finally pass this
# list to tf.keras.layers.Add() to sum the values element-wise.
net = tf_agents.networks.Sequence([
  NestMap({"inp1": layer1, "inp2": layer2}),
  NestFlatten(),
  tf.keras.layers.Add(),
])
combined_outputs, next_state = net({"inp1": inp1, "inp2": inp2}, state)