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