python - How to dynamically set an object's method at runtime (when the object is an instance of a Gym environment)?

Question

I want to override an instance's method. Let's call this instance/object env. The instance method has a method with the signature f(self, other_parameter), which is actually not publically accessible (not sure why, but you will see below that it's not listed in the list of properties of the class and object: can you please tell me why?). However, I'm given functions of the form new_f(other_parameter), i.e. without the self parameter, which I need to bind to the instance/object, according to this answer. However, this is not my first problem.

My first problem is that no method that I've tried from this post seems to work in my case, i.e. the new method is not assigned to the object: at least, this is what I deduce from the fact that the method doesn't seem to be called. You can run the following example to see the problem.

import types
from pprint import pprint

import gym


class MyEnv(gym.Env):
    def f(self, other_parameter):
        print("other_parameter (original) =", other_parameter)
        return 0

    def step(self, action: int) -> tuple:
        return [], self.f(action), False, {}

    def reset(self):
        return []


env_id = "my_env-v1"
env = MyEnv()
gym.envs.register(id=env_id,
                  entry_point=lambda: env,
                  max_episode_steps=200)

env = gym.make(env_id)

pprint(dir(env))

env.reset()

# CORRECTLY prints the following
# other_parameter (original) = 0
# ([], 0, False, {})
print(env.step(0))

print("-" * 10)


def f(self, other_parameter):
    print("other_parameter (new) =", other_parameter)
    return 10


# IT DOES NOT SEEM TO TAKE ANY EFFECT.
# env.f = f

# EVEN IF I TRY OTHER METHODS, IT DOES NOT SEEM TO WORK.
env.f = types.MethodType(f, env)

# INCORRECTLY prints the following
# other_parameter (original) = 7
# ([], 0, False, {})
# IT SHOULD PRINT
# other_parameter (new) = 7
# ([], 10, False, {})
print(env.step(7))

I'm using gym 0.17.3 (just do pip install gym==0.17.3 to install it).

This problem probably occurs because Gym environments are created by calling gym.make(env_id), which probably does some magic or weird tricks. I've already looked at this module where make is defined, but I'm not yet sure where the problem could be. I really need to set the method of a created environment once it's been created, so if you have a suggestion of how to do it, I would appreciate it.

Answer 1

I've just noticed that if I try to bind the new method f to env before the line gym.envs.register(...), it works, but I would like to change f after the registration.

Meanwhile, I've also noticed that gym can wrap your environment object into another class. In the example above, it turns out that an MyEnv object is created inside an TimeLimit object, which is what env variable in the example above points to. I also notice that env has another property called env, which is the actual object of type MyEnv, so my example above works by only changing the line

env.f = types.MethodType(f, env)

to

env.env.f = types.MethodType(f, env)

This also explains why before registering the environment it works, because env before registering the environment is actually an instance of MyEnv.