Discussion of Vector Environment API Breaking Changes

[jkterry1]

There are numerous long standing complaints about the vector environment API, to the point where it is the one thing that most people have seemed to think breaking changes are warranted for (which I agree with). I've created this thread to come to a consensus on the solution.

Because whatever finalized changes we come up with here will be very significant to the community and this can only really be done once, I am specifically insisting that the following people agree to whatever breaking changes are made:

Benjamin Black (@benblack769, UMD and PettingZoo)
Antonin Raffin (@araffin, DLR RM and Stable Baselines)
Chris Nota (@cpnota, UMass Amherst and the Autonomous Learning Library)
Costa Huang (@vwxyzjn, Weights and Biases and CleanRL)
Jesse Fareboro (@JesseFarebro, MILA and the Arcade Learning Environment)

[joschu]

You might want to check out https://github.com/openai/gym3

[tristandeleu]

There are numerous long standard complaints about the vector environment API, to the point where it is the one thing that most people have seemed to think breaking changes are warranted for (which I agree with).

Can you be more precise about these complaints that warrant breaking the API, in order to get a better understanding of what needs to be done? As far as I know, almost all the issues involving VectorEnv here have been addressed one way or another.

To be specific, there are a number of things that can be improved in my opinion, but very few of them require breaking the API:

• Add documentation. This should be the number 1 priority, but unfortunately it does not seem to be part of the current plan for the documentation Specific Plans for Documentation Website #2276
• Refactor the _worker and _worker_shared_memory in order to avoid code duplication and to simplify custom workers (custom workers should be treated as an advanced feature anyway)
• Add call/get_attr/set_attr Add call, get_attr and set_attr methods to VectorEnv #1600
• Allow nested VectorEnv, e.g. SyncVectorEnv inside a AsyncVectorEnv, to address Chunking in vector environments #1795 & [gym.vector] Add BatchedVectorEnv, (chunking + flexible n_envs) #2072. This is probably done in Bugfix: Allow Nesting of Sync/Async VectorEnvs #2104.
• Renaming AsyncVectorEnv (and possibly SyncVectorEnv) to something more meaningful, since this is not the process is not asynchronous but parallel. This was a poorly chosen naming convention I introduced to match the names in Ray/RLlib, but it should be something like ParallelVectorEnv and SequentialVectorEnv. This would be the only breaking change to the API.

You might want to check out https://github.com/openai/gym3

The problem with Gym3 is that it uses a different API than Gym (obs_space instead of observation_space, observe instead of step, even the spaces are different e.g. the Tuple space does not exist), so anything related to Gym3 requires a wrapper FromGymEnv to convert from a Gym environment to a Gym3 one. I also believe (I might be wrong though) that anything Gym3's SuprocEnv can do, Gym's AsyncVectorEnv can as well and more (e.g. using shared memory for observations, similar to baselines' ShmemVecEnv).

[benblack769]

My main complaint is the tuple observation/action spaces. Stable Baseline's vector environment has the observation/space just be the observation space of one of the environments, and assumes that all environments have the same spaces. This seems much more sensible to me than having a tuple observation space because the main use case of a vector environment is to put many near-identical environments together, and applying a single policy to many of them. Having a single observation space makes using the vector environment in this way much easier. Especially when you start writing vector environment wrappers, etc.

While I agree that the tuple space is more general, in this case I prefer simplicity over generality.

[benblack769]

In general, I would prefer it if Gym adopted Stable Baselines vector environment API.

[vwxyzjn]

I would like to second @benblack769: The design of the action space of the vectorized env is cleaner in SB3 in my opinion, and I would prefer if gym adopted the SB3’ vector environment API.

It seems the only utility of the tuple action space is so that you can do env.step(env.action_space.sample()), which I think it’s ok to deprecate.

The two biggest benefits of gym3 API appears to be:

1. Implementing multi-agent environment becomes very straightforward (see here).
2. Its API design of act() and observe() allow the env to step asynchronously.

Regarding the first benefit, you could use SB3’s API as well (see here as an example).

Regarding the second benefit, I feel it was set up for more complicated use cases (e.g. distributed training, asynchronous rollouts), so I am inclined to keep it simple and just use SB3’s API.

[benblack769]

I would note that implementing a multi-agent environment is equally simple using SB3's vector env API. For example: https://github.com/PettingZoo-Team/SuperSuit/#parallel-environment-vectorization

[tristandeleu]

My main complaint is the tuple observation/action spaces. Stable Baseline's vector environment has the observation/space just be the observation space of one of the environments, and assumes that all environments have the same spaces. This seems much more sensible to me than having a tuple observation space because the main use case of a vector environment is to put many near-identical environments together, and applying a single policy to many of them.

I'm confused, Gym's VectorEnv does not return a tuple of observations. For example:

import gym
env = gym.vector.make('CartPole-v1', num_envs=5)
observations = env.reset()
# array([[-0.04071225, -0.00272771,  0.04585911,  0.03397448],
#        [ 0.04807664, -0.00466015,  0.03836329, -0.04897803],
#        [ 0.03999503,  0.03953767, -0.02474664, -0.00865895],
#        [-0.03199483,  0.01612429, -0.03704762, -0.02390875],
#        [ 0.00967298,  0.01544605,  0.04391582,  0.0040252 ]],
#       dtype=float32)

And this is true for any (arbitrarily nested) observation space, see #1513 for additional examples (in particular, one with a Dict observation space).

The only instance where it will return a tuple of observations is whenever the observation space is not a standard Gym space (e.g. Box, Discrete, Dict, etc...), because in that case batching would depend on the nature of the data. This is a feature introduced here #2038 (see the corresponding issues in that PR for the use cases).

In general, I would prefer it if Gym adopted Stable Baselines vector environment API.

As far as I know, Gym's VectorEnv and SB3's VecEnv APIs are almost identical, because both were created on top of baseline's SubprocVec.

The one difference I can spot is that Gym's VectorEnv inherits from gym.Env, whereas SB3's VecEnv does not. This means that VectorEnv needs to implement both observation_space and action_space, and should be such that, for example, an observation from VectorEnv must be a valid element of observation_space, which is why VectorEnv's observation_space is not the observation space of a child environment (unlike in SB3's VecEnv).

Here, there is indeed a question of wether to keep the action_space as a tuple of individual spaces, or having it batched as well. Note that other than this convention for action_space (which should only matter if you have complex Tuple/Dict action spaces), under the hood Gym's AsyncVectorEnv and SB3's SubprocVecEnv have the exact same implementation.

The design of the action space of the vectorized env is cleaner in SB3 in my opinion, and I would prefer if gym adopted the SB3’ vector environment API.

I don't understand what you mean by the action space being cleaner? If it is because you can feed a numpy array directly in step_async/step, then you can do exactly the same in Gym (again, the implementation is really the same in both cases). This is valid as long as you can iterate over the action fed to step_async/step (SB3 is making the same assumption)

import gym
import numpy as np

env = gym.vector.make('MountainCarContinuous-v0', num_envs=5)
observations = env.reset()

# Here actions.shape = (5, 1) because the action space
# of MountainCarContinuous-v0 has shape (1,)
actions = np.array([[-1.], [-0.5], [0.], [0.5], [1.]])
observations, rewards, dones, infos = env.step(actions)

Regarding the first benefit, you could use SB3’s API as well (see here as an example).

This is also available in Gym (and subclasses, e.g. in AsyncVectorEnv).

I feel like a lot of the misunderstanding could be cleared up with proper documentation.

[benblack769]

@tristandeleu

Compare:

>>> import gym
>>>
>>> import numpy as np
>>>
>>>
>>>
>>> env = gym.vector.make('MountainCarContinuous-v0', num_envs=5)
>>> env.action_space
Tuple(Box(-1.0, 1.0, (1,), float32), Box(-1.0, 1.0, (1,), float32), Box(-1.0, 1.0, (1,), float32), Box(-1.0, 1.0, (1,), float32), Box(-1.0, 1.0, (1,), float32))
>>> env.observation_space
Box(-1.2000000476837158, 0.6000000238418579, (5, 2), float32)
>>>

To:

>>> import stable_baselines3
>>> vec_env = stable_baselines3.common.vec_env.DummyVecEnv([lambda: gym.make("MountainCarContinuous-v0")]*5)
>>> vec_env.action_space
Box(-1.0, 1.0, (1,), float32)
>>> vec_env.observation_space
Box(-1.2000000476837158, 0.6000000238418579, (2,), float32)

Please reread the above complaints with this in mind.

[tristandeleu]

Just to emphasize that the action_space as a Tuple or a batch in VectorEnv is really just a matter of convention, if you batch the action space (see the diff below), then all the tests are still passing.

diff --git a/gym/vector/vector_env.py b/gym/vector/vector_env.py
index 375826f..7d09815 100644
--- a/gym/vector/vector_env.py
+++ b/gym/vector/vector_env.py
@@ -33,7 +33,7 @@ class VectorEnv(gym.Env):
         super(VectorEnv, self).__init__()
         self.num_envs = num_envs
         self.observation_space = batch_space(observation_space, n=num_envs)
-        self.action_space = Tuple((action_space,) * num_envs)
+        self.action_space = batch_space(action_space, n=num_envs)

         self.closed = False
         self.viewer = None

Please reread the above complaints with this in mind.

That was my point in the comparison between Gym and SB3: in SB3's VecEnv, an action you feed to step is not a valid element of action_space (which is ok, since VecEnv does not inherit from gym.Env, so action_space and observation_space may have different semantics). In that example, you will feed a numpy array of shape (5, 1) to DummyVecEnv, not a numpy array of shape (1,). In Gym's VectorEnv (just like any gym.Env), under the convention of having a batched space (i.e. the diff above is applied), the action is going to be a valid element of action space, since:

# With the diff above
env = gym.vector.make('MountainCarContinuous-v0', num_envs=5)
print(env.action_space)
# Box(-1.0, 1.0, (5, 1), float32)

[cpnota]

The issue I (or actually @benblack769, who made the changes) had with the Gym VectorEnv when using it in the autonomous-learning-library is that it does not return the observation associated with terminal states, unlike the standard API. Instead, the done flag is associated with the first observation of the next episode. This created some inconsistencies in ParallelAgent API. I would prefer VectorEnv to return the observations and rewards associated with the terminal states separately from the first state of the next episode.

[cpnota]

On the subject of action_space:

>>> vec_env.action_space
Box(-1.0, 1.0, (1,), float32)
>>> vec_env.observation_space
Box(-1.2000000476837158, 0.6000000238418579, (2,), float32)

I agree this is a pretty ugly inconsistency. I see @tristandeleu just opened a PR to address this.

[tristandeleu]

The issue I (or actually @benblack769, who made the changes) had with the Gym VectorEnv when using it in the autonomous-learning-library is that it does not return the observation associated with terminal states, unlike the standard API. Instead, the done flag is associated with the first observation of the next episode.

That's a good point, I agree this is not ideal. Originally I had a key in the info dictionary with the last observation #1513 (comment) (it was tied to another feature dropped though), but I removed it altogether to match baselines' implementation. I have nothing against adding a key to the info dict again, similar to how SB3 does it.

[vwxyzjn]

@tristandeleu you have a valid point. The action with shape (1,) is indeed not a valid action to step, and I think using the batch_space would make the API better!

That said, jus throwing some nit thoughts about the usability. Currently, I have been writing code like the following with SB3's API:

self.actor = nn.Sequential(
    layer_init(nn.Linear(np.array(envs.observation_space.shape).prod(), 64)),
    nn.Tanh(),
    layer_init(nn.Linear(64, 64)),
    nn.Tanh(),
    layer_init(nn.Linear(64,  envs.action_space.n), std=0.01),
)

I can do things like nn.Linear(np.array(envs.observation_space.shape).prod(), 64) because the "true" observation space stays the same regardless of the number of environments or the number of samples used during training.

Using the new batch_spaceed API, I imagine I would need to refactor it to something like nn.Linear(np.array(envs.observation_space.shape[1:]).prod(), 64) Would this be confusing?

We would need to clarify that the first dimension of observation_space's shape is a "batch" dimension and therefore should be ignored during the NN constructions. Maybe with more documentation then this is a non-issue.

[tristandeleu]

@vwxyzjn the properties single_observation_space and single_action_space got you covered!
Depending on how much flexibility you need, you can also combine that with the flatdim utility function. Something like

self.actor = nn.Sequential(
    layer_init(nn.Linear(flatdim(envs.single_observation_space), 64)),
    nn.Tanh(),
    layer_init(nn.Linear(64, 64)),
    nn.Tanh(),
    layer_init(nn.Linear(64, envs.single_action_space.n), std=0.01),
)