def unroll_fn(inputs: types.Observation,
state: hk.LSTMState) -> HaikuLSTMOutputs:
model = networks_lib.DeepIMPALAAtariNetwork(
env_spec.actions.num_values)
return model.unroll(inputs, state)
def initial_state_fn(batch_size: Optional[int] = None) -> hk.LSTMState:
model = networks_lib.DeepIMPALAAtariNetwork(
env_spec.actions.num_values)
return model.initial_state(batch_size)