Call batch_concat to support nested inputs for discrete actions.
@@ -244,7 +244,8 @@ def make_discrete_networks(
layers.extend([networks_lib.AtariTorso()])
layers.extend([hk.nets.MLP(hidden_layer_sizes, activate_final=True)])
trunk = hk.Sequential(layers)
- h = trunk(inputs)
+ h = utils.batch_concat(inputs)
+ h = trunk(h)
logits = hk.Linear(num_actions)(h)
values = hk.Linear(1)(h)
values = jnp.squeeze(values, axis=-1)
