[Update core.py]: support Wrapper for registering global methods with nested wrapping

[zuoxingdong]

According to the desired functionalities in #1319, I would like to propose to support Wrapper class for registering arbitrary methods which can be accessible to the entire nested wrapping.

e.g.

class A(Wrapper):
    def __init__(self, env):
        super().__init__(env)
        
        self._register('action_flatdim', property(lambda self: self.action_space.n))
        self._register('ONE', 1)
        
        
class B(Wrapper):
    def __init__(self, env):
        super().__init__(env)
        
        self.action_space = spaces.Discrete(50)
        
        self._register('TWO', 2)

env = gym.make('CartPole-v1')
env = A(env)
env = B(env)

print(env.env.action_flatdim)
print(env.action_flatdim)
2
50

print(env.ONE)
1

What do you think ? @pzhokhov @christopherhesse

[pzhokhov]

this logic means that child classes have to call "_register" before the call to Wrapper.init superconstructor, correct?

[pzhokhov]

for instance, if the example you are providing, if you add "print(env.TWO)" it will throw an error

[pzhokhov]

@zuoxingdong I like the idea (we had discussions around similar ideas internally not long ago - mentioning relevant people here: @christopherhesse @kcobbe). However, there is one implementation issue (mentioned in-code) and also as issue of the possible name collisions. One possible solution to name collisions is to throw a warning and return most-recent result (which is what your code would do currently, from what I understand).

[christopherhesse]

I think having some way of searching through wrappers for a method besides the normal gym env methods would be nice, but it's not clear what the least bad way of doing this is.

I was thinking something along the lines of:

class Wrapper:
    def call(self, name, *args, **kwargs):
        if name.startswith('_'):
            return None
    
        if hasattr(self, name):
            return getattr(self, name)(*args, **kwargs)
        elif self.env.unwrapped is self.env:
            # reached underlying gym environment
            return None
        else:
            return self.env.unwrapped.call(name, *args, **kwargs)
        

class A(Wrapper):
    def __init__(self, env):
        super().__init__(env)

    def get_info(self):
        info = self.env.call('get_info') or {}
        return {**info, 'a': 1}

    def _internal_method(self):
        return self.action_space.n

        
class B(Wrapper):
    def __init__(self, env):
        super().__init__(env)
        
        self.action_space = spaces.Discrete(50)

    def get_info(self):
        info = self.env.call('get_info') or {}
        return {**info, 'b': 2}
        

env = gym.make('CartPole-v1')
env = A(env)
env = B(env)

print(env.call('get_info'))

This avoids using the very magical __getattr__ stuff. I think this would actually require that call() is part of the Gym Env api as well, this could be implemented externally to the API as a separate function like:

def call_method(env, name):
      # search through all wrappers for a specific method name

Any other suggestions?

[zuoxingdong]

@pzhokhov @christopherhesse That sounds a great idea ! A tiny follow-up, how about to unify the search for both attributes and methods, e.g.

class Wrapper(object):
    def call(self, name):
        if name.startswith('_'):
            raise ValueError('not allowed to search for private members.')
        if hasattr(self, name):
            return getattr(self, name)
        elif self.env.unwrapped is self.env:
            # reached underlying gym environment
            return None
        else:
            return self.env.unwrapped.call(name)

[christopherhesse]

I don't like my previous suggestion very much. I think one problem with existing wrappers is that Wrappers only wrap the Gym API, not just arbitrary python objects. So if your environment has a get_state() method and you wrap it, you can't access it, except through env.unwrapped which prevents wrappers from intercepting that method call.

I will type up an example that might be better than my previous proposal

[christopherhesse]

So if wrappers actually wrap the object in the normal python sense, then they should probably forward all attribute lookups:

import gym
import gym.spaces
import numpy as np

# existing wrapper

class StateEnv(gym.Env):
    def __init__(self):
        self.state = 1
        self.action_space = gym.spaces.Discrete(4)

    def reset(self):
        return None

    def step(self, act):
        return None, 0.0, False, {}

    def get_state(self):
        return self.state


class ClipRewards(gym.Wrapper):
    def step(self, act):
        obs, rew, done, info = self.env.step(act)
        return obs, np.sign(rew), done, info


env = StateEnv()
env.get_state()
env = ClipRewards(env)
# env.get_state() <- doesn't work
env.unwrapped.get_state()

# wrapper that forwards getattr

class Wrapper2(gym.Wrapper):
    def __getattr__(self, name):
        if name.startswith('_'):
            raise AttributeError("attempted to get missing private attribute '{}'".format(name))
        return getattr(self.env, name)


class ClipRewards2(Wrapper2):
    def step(self, act):
        obs, rew, done, info = self.env.step(act)
        return obs, np.sign(rew), done, info

env = StateEnv()
env.get_state()
env = ClipRewards2(env)
env.get_state() # <- works

Adding methods via wrapper seems fine here too:

class FlatDim(Wrapper2):
    def action_flatdim(self):
        return self.action_space.n

env = StateEnv()
# env.action_flatdim() # <- this won't work, but probably it shouldn't
env = FlatDim(env)
env.action_flatdim()

[christopherhesse]

This seems fine for calls that we want to intercept in the wrappers as well:

class GetLogInfo1(Wrapper2):
    def get_log_info(self):
        log_info = {}
        if hasattr(self.env, 'get_log_info'):
            log_info = self.env.get_log_info()
        return {**log_info, 'a': 1}

class GetLogInfo2(Wrapper2):
    def get_log_info(self):
        log_info = {}
        if hasattr(self.env, 'get_log_info'):
            log_info = self.env.get_log_info()
        return {**log_info, 'b': 2}

env = StateEnv()
# env.get_log_info() <- this would fail, caller must use hasattr() and either throw an exception because the method is required or else provide a default value
env = GetLogInfo1(env)
print(env.get_log_info())
env = GetLogInfo2(env)
print(env.get_log_info())

[christopherhesse]

What do you think @zuoxingdong and @pzhokhov ?

[zuoxingdong]

@christopherhesse It looks great, I like this clean solution which allows to search arbitrary python objects (both attributes and methods) with hierarchical priority from the root to the leaf.

Sometimes it can also be convenient to access objects in specific Wrapper, this can be supported by a simple external API, e.g.

def get_wrapper(env, name):
    if name == env.__class__.__name__:
        return env
    elif env.unwrapped is env:  # reaching underlying environment
        return None
    else:
        return get_wrapper(env.env, name)

[christopherhesse]

I can't think of a use case for accessing specific wrappers, do you have one?

[zuoxingdong]

I am thinking of the use case of observation/reward standardization wrapper, during evaluation phase, we might want to retrieve the current estimation of mean and std to apply Z-filter to evaluation environment. When there are other wrappers on top of standardization wrapper, it is still true that Wrapper2 implementation can easily take the moments objects directly, but it can fail to work when other wrapper on top of it has python object defined with same name of estimated moments. Well, I'm hesitating if this is really necessary.

[christopherhesse]

Yeah, I think I'd prefer to wait until that actually happened, since you can just call your property/method whatever you want to avoid conflicts with other wrappers.

[zuoxingdong]

I agree, also, the user can easily handle these cases once it happens. For now, how about let's ship the current __getattr__ functionality, what do you think ?

[christopherhesse]

Sure, if you wanted to update your PR to change gym.Wrapper to use the __getattr__ setup, @pzhokhov can review. Thanks for your help on this!