HumanCompatibleAI / imitation

Clean PyTorch implementations of imitation and reward learning algorithms
https://imitation.readthedocs.io/
MIT License
1.33k stars 246 forks source link

Problem in Passing Along Custom Gym Env Constructor Parameters in make_vec_env #850

Open christianjcc opened 6 months ago

christianjcc commented 6 months ago

Bug description

I was testing out the imitation learning library with a custom gym environment and ran into a shortcoming in imitation/util/util.py. I get the error message provided below.

`Traceback (most recent call last): File "/home/anaconda3/envs/env/lib/python3.8/site-packages/gymnasium/envs/registration.py", line 802, in make env = env_creator(**env_spec_kwargs) TypeError: init() missing 1 required positional argument: 'config'

During handling of the above exception, another exception occurred:

Traceback (most recent call last): File "runner.py", line 118, in main() File "runner.py", line 58, in main env = make_vec_env( File "/home/anaconda3/envs/educagent_libtraffic/lib/python3.8/site-packages/imitation/util/util.py", line 117, in make_vec_env tmp_env = gym.make(env_name) File "/home/anaconda3/envs/env/lib/python3.8/site-packages/gymnasium/envs/registration.py", line 814, in make raise type(e)( TypeError: init() missing 1 required positional argument: 'config' was raised from the environment creator for custom/mycustom-env with kwargs ({}) `

Below is an example of how I define, register the custom gym environment, and pass along the config dictionary to the env instance, `

Define your config dictionary

config = { "parameters": "example", "render_mode": "human", }

env_name="custom/mycustom-env" gym.register( id=env_name, entry_point=MyCustomEnv, max_episode_steps=500, )

env = make_vec_env(
    env_name=env_name,
    n_envs=1,
    post_wrappers=[
        lambda env, _: RolloutInfoWrapper(env)
    ],  # needed for computing rollouts later
    env_make_kwargs={"config": config},
)

I can get it to work with the way the imitation library has implemented it's gym.make instance by doing the following instead:

Define a factory function that returns the custom environment class with the desired configuration

def make_traffic_env(**kwargs):
    return MyCustomEnv(config=config, **kwargs)

env_name="custom/mycustom-env"

gym.register(
    id=env_name,
    entry_point=Tmake_traffic_env,
    max_episode_steps=500,
)

env = make_vec_env(
    env_name=env_name,
    n_envs=1,
    post_wrappers=[
        lambda env, _: RolloutInfoWrapper(env)
    ],  # needed for computing rollouts later
)

`

This occurs because the following variable is defined as follows: tmp_env = gym.make(env_name)

As implemented here: https://github.com/HumanCompatibleAI/imitation/blob/master/src/imitation/util/util.py#L117

instead of: tmp_env = gym.make(env_name, **env_make_kwargs)

to pass along the kwargs as defined in, https://gymnasium.farama.org/api/registry/#gymnasium.make

By implementing the suggestion above, it will help users avoid having to define factory function, and simplifying the steps.

Environment

chrisjcc commented 5 months ago

according to https://imitation.readthedocs.io/en/latest/, imitation is built on and compatible with Stable Baselines 3 (SB3), would it have trouble running with Ray RLlib, or are there any plans to support imitation with RLlib?