Sample action from an action space for given state

def step(self, observations):
        """ Sample action from an action space for given state """
        q_values = self(observations)

        return {
            'actions': self.q_head.sample(q_values),
            'q': q_values
        }