def q(obs):
flatten = lambda x: jnp.reshape(x, (-1, ))
network = hk.Sequential(
[flatten, nets.MLP([FLAGS.hidden_units, num_actions])])
return network(obs)
def q(obs):
network = hk.Sequential(
[hk.Flatten(),
nets.MLP([FLAGS.hidden_units, num_actions])])
return network(obs)
