TFEnvironment not calling custom PyEnvironment overrides

General Discussion
,
#1

I have a custom PyEnvironment which overrides the necessary functions.

def action_spec(self)
def observation_spec(self)
def _reset(self)
def _step(self, action)

I instantiate the environment and validate it works just fine
env = MyEnv()
utils.validate_py_environment(env, episodes=5)

I create a TFPyEnvironment wrapper for my environment
tf_env = tf_py_environment.TFPyEnvironment(env)

…and when I call reset() on tf_env, it does not call the reset function of MyEnv. To my understanding, the TFPyEnvironment reset() should call reset() on the PyEnvironment it wraps.

tf_env.reset() # does not call the .reset() from MyEnv

I am running TensorFlow 2.6.0 and tf_agents 0.9.0