code: def step(self, observation, argmax_sampling=False):
          
          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
          }
docs: Select actions based on model's output