def get_action(self, observation):
    'Return a single action.'
    flat_obs = self.observation_space.flatten(observation)
    prob = self._f_prob([flat_obs])[0]
    action = self.action_space.weighted_sample(prob)
    return (action, dict(prob=prob))
Return a single action.
def get_action(self, observation):
    
    flat_obs = self.observation_space.flatten(observation)
    prob = self._f_prob([flat_obs])[0]
    action = self.action_space.weighted_sample(prob)
    return (action, dict(prob=prob))