Select actions based on model's output

def step(self, observation, argmax_sampling=False):
        """ Select actions based on model's output """
        policy_params, q = self(observation)
        actions = self.action_head.sample(policy_params, argmax_sampling=argmax_sampling)

        # log probability - we can do that, because we support only discrete action spaces
        logprobs = self.action_head.logprob(actions, policy_params)

        return {
            'actions': actions,
            'q': q,
            'logprobs': policy_params,
            'action:logprobs': logprobs
        }