Custom Exploration Behavior based on Observations


I am trying to write a custom exploration strategy that decides which explorative action to take based on the observations. I think I need to implement a custom Exploration class. However, It seems the observations are not passed to the get_exploration_action function. Any idea how I can also get the observation inside this function?

def get_exploration_action(self, *,
       action_distribution: ActionDistribution,
       timestep: Union[TensorType, int],
       explore: bool = True):