def __init__(self, num_actions: int):
super().__init__(name='impala_atari_network')
self._embed = embedding.OAREmbedding(
DeepAtariTorso(use_layer_norm=True), num_actions)
self._core = hk.GRU(256)
self._head = policy_value.PolicyValueHead(num_actions)
self._num_actions = num_actions
def __init__(self, num_actions: int):
super().__init__(name='r2d2_atari_network')
self._embed = embedding.OAREmbedding(DeepAtariTorso(), num_actions)
self._core = hk.LSTM(512)
self._duelling_head = duelling.DuellingMLP(num_actions,
hidden_sizes=[512])
self._num_actions = num_actions