rllib use custom registered environments

Question

Rllib docs provide some information about how to create and train a custom environment. There is some information about registering that environment, but I guess it needs to work differently than gym registration.

I'm testing this out working with the SimpleCorridor environment. If I add the registration code to the file like so:

from ray.tune.registry import register_env

class SimpleCorridor(gym.Env):
   ...


def env_creator(env_config):
    return SimpleCorridor(env_config)

register_env("corridor", env_creator)

Then I am able to train an algorithm using the string name no problem:

if __name__ == "__main__":
    ray.init()
    tune.run(
        "PPO",
        stop={
            "timesteps_total": 10000,
        },
        config={
            "env": "corridor", # <--- This works fine!
            "env_config": {
                "corridor_length": 5,
            },
        },
    )

However

It is kinda pointless to register the environment in the same file that you define the environment because you can just use the class. OpenAI gym registration is nice because if you install the environment, then you can use it anywhere just by writing

include gym_corridor

It's not clear to me if there is a way to do the same thing for registering environments for rllib. Is there a way to do this?

score 11 · Accepted Answer · answered Mar 21 '20 at 20:22

The registry functions in ray are a massive headache; I don't know why they can't recognize other environments like OpenAI Gym.

Anyway, the way I've solved this is by wrapping my custom environments in another function that imports the environment automatically so I can re-use code. For example:

def env_creator(env_name):
    if env_name == 'CustomEnv-v0':
        from custom_gym.envs.custom_env import CustomEnv0 as env
    elif env_name == 'CustomEnv-v1':
        from custom_gym.envs.custom_env import CustomEnv1 as env
    else:
        raise NotImplementedError
    return env

Then, to get it to work with the tune.register_env(), you can use your custom env with a lambda function:

env = env_creator('CustomEnv-v0')
tune.register_env('myEnv', lambda: config, env(config))

From there, tune.run() should work. It's annoying, but that's the best way I've found to work around this registry issue.

Yes, this is what I've ended up doing. Still not idea, as you say, but it works. — KindaTechy, Mar 23 '20 at 00:34

score 1 · Answer 2 · answered Aug 12 '20 at 15:58

Here's an example of defining a Gym custom environment and registering it for use in both Gym and RLlib https://github.com/DerwenAI/gym_example

See the Python example code in:

sample.py – how to create an agent using gym.make() to measure the performance of a random-action baseline
train.py – register, train a policy with RLlib, measure performance of learning, show a programmatic rollout

The file structure of the environment's Git repo is tricky, but that allows for a Python import of the environment from a Git repo, pip, conda, etc. -- related to what you were asking. I agree that the SimpleCorridor example is almost pointless since it registers and uses a custom environment in the same file that defines the environment's class. Likewise, that example shows how to measure learning with RLlib, but fails to show how a policy could ever be used, i.e., how to restore and deploy a checkpoint of a trained policy in a use case. An upcoming blog post for Ray explores gym_example in more detail.

Here's an article that describes the repo above in more detail https://medium.com/distributed-computing-with-ray/anatomy-of-a-custom-environment-for-rllib-327157f269e5 — Paco, Aug 31 '20 at 23:01

rllib use custom registered environments

2 Answers2

Linked