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?
Answers:
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.
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.
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?
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.
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 usinggym.make()
to measure the performance of a random-action baselinetrain.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.