Select actions based on model's output

def step(self, observations, state, argmax_sampling=False):
        """ Select actions based on model's output """
        action_pd_params, value_output, new_state = self(observations, state)
        actions = self.action_head.sample(action_pd_params, argmax_sampling=argmax_sampling)

        # log likelihood of selected action
        logprobs = self.action_head.logprob(actions, action_pd_params)

        return {
            'actions': actions,
            'values': value_output,
            'logprobs': logprobs,
            'state': new_state
        }