kengz/SLM-Lab

VectorEnvWrapper base missing spec

DuaneNielsen opened this issue · 2 comments

Describe the bug
when writing a new VectorEnvWrapper by extending baseclass VectorEnvWrapper, it appears calling super is not enough to make it work. Seems that the OpenAIEnv constructor expects a wrapper to carry a spec attribute, which is not present in VectorEnvWrapper base class.

To Reproduce

  1. OS and environment: Ubuntu 18.04, conda, etc as per docs
  2. SLM Lab git SHA (run git rev-parse HEAD to get it): 41e6918
  3. spec file used: not relevant, self evident from code review

Additional context

The below is required to implement a working wrapper... removing the self.spec = venv.spec causes it to fail, due to the reference in the OpenAIEnv class constructor.

class VectorLogWrapper(VecEnvWrapper):
    def __init__(self, venv, arg_one, arg_two):
        self.spec = venv.spec
        super().__init__(venv, venv.observation_space, venv.action_space)
        self.arg_one = arg_one
        self.arg_two = arg_two

    def step_wait(self):
        obs, rews, news, infos = self.venv.step_wait()
        logger.info(f'step_called arg_one: {self.arg_one}')
        return obs, rews, news, infos

    def reset(self):
        s = self.venv.reset()
        logger.info('reset_called')
        return s

The reference to spec is used in OpenAIEnv on the below line.

self.max_t = self.max_t or self.u_env.spec.max_episode_steps

Context below...

class OpenAIEnv(BaseEnv):
    '''
    Wrapper for OpenAI Gym env to work with the Lab.

    e.g. env_spec
    "env": [{
        "name": "PongNoFrameskip-v4",
        "frame_op": "concat",
        "frame_op_len": 4,
        "normalize_state": false,
        "reward_scale": "sign",
        "num_envs": 8,
        "max_t": null,
        "max_frame": 1e7
    }],
    '''

    def __init__(self, spec):
        super().__init__(spec)
        try_register_env(spec)  # register if it's a custom gym env
        seed = ps.get(spec, 'meta.random_seed')
        episode_life = not util.in_eval_lab_modes()
        if self.is_venv:  # make vector environment
            self.u_env = make_gym_venv(spec=spec['env'][0], name=self.name, num_envs=self.num_envs, seed=seed, frame_op=self.frame_op, frame_op_len=self.frame_op_len, image_downsize=self.image_downsize, reward_scale=self.reward_scale, normalize_state=self.normalize_state, episode_life=episode_life)
        else:
            self.u_env = make_gym_env(spec=spec['env'][0], name=self.name, seed=seed, frame_op=self.frame_op, frame_op_len=self.frame_op_len, image_downsize=self.image_downsize, reward_scale=self.reward_scale, normalize_state=self.normalize_state, episode_life=episode_life)
        if self.name.startswith('Unity'):
            # Unity is always initialized as singleton gym env, but the Unity runtime can be vec_env
            self.num_envs = self.u_env.num_envs
            # update variables dependent on num_envs
            self._infer_venv_attr()
            self._set_clock()
        self._set_attr_from_u_env(self.u_env)
        self.max_t = self.max_t or self.u_env.spec.max_episode_steps
        assert self.max_t is not None
        logger.info(util.self_desc(self))

Perhaps this should be initialized in the base class? Not sure of the correct fix as I"m not yet familiar with the overall design.

kengz commented

Hi, this specific spec object is required by the OpenAI gym interface, and is an instance of its EnvSpec class. That said, any environment/wrapper extending from OpenAI gym interface should follow its API and include the spec object.

Additionally, a Wrapper class is distinct from an Env class, and it needs to inherit that from its env object. You may follow this example to do so. Let me know if this solves your problem!

Thanks. Yes that's the example I was following and that's the way I "fixed" the problem.

On the face of it though, it just looks like there is problem in the inheritance hierarchy. It's unusual to have to add a "magic" value to an base class function to make things work, even if it's an "edge case. Usually the convention is that the base class constructor requires you to add everything needed.

Anyway, all good, just thought I'd mention it in case it was an oversight. Seems like that's the design, which is OK.