@callback init(random_key :: Nx.t(), opts :: keyword()) :: {t(), random_key :: Nx.t()}

Initializes the agent state with the given agent-specific options.

Should be implemented in a way that the result would be semantically the same as if reset/2 was called in the end of the function.

As a suggestion, the implementation should only initialize fixed values here, that is values that don't change between sessions (epochs for non-episodic tasks, episodes for episodic tasks). Then, call reset/2 internally to initialize the rest of variable values.

@callback record_observation(
  rl_state(),
  action :: Nx.t(),
  reward :: Nx.t(),
  is_terminal :: Nx.t(),
  next_rl_state :: rl_state()
) :: rl_state()

Can be used to record the observation in an experience replay buffer.

If this is not desired, just make this function return the first argument unchanged.

@callback reset(random_key :: Nx.t(), rl_state :: t()) :: {t(), random_key :: Nx.t()}

Resets any values that vary between sessions (which would be episodes for episodic tasks) for the agent state.

@callback select_action(rl_state(), iteration :: Nx.t()) :: {action :: Nx.t(), rl_state()}

Selects the action to be taken.