code: def step(self, observations):
          
          q_values = self(observations)
  
          return {
              'actions': self.q_head.sample(q_values),
              'q': q_values
          }
docs: Sample action from an action space for given state