Batch action_space in VectorEnv

[tristandeleu]

Given the discussion in #2279, here is a proposal to have a batch action_space instead of a Tuple instance.

import gym
env = gym.vector.make('CartPole-v1', num_envs=5)
observations = env.reset()

print(env.action_space)
# Before: Tuple(Discrete(2), Discrete(2), Discrete(2), Discrete(2), Discrete(2))
# After: MultiDiscrete([2 2 2 2 2])

actions = np.array([1, 0, 0, 1, 1])
observations, rewards, dones, infos = env.step(actions)
print(f'Observations shape: {observations.shape}')
# Observations shape: (5, 4)

This handles any nested action space as well:

import gym
import numpy as np

from gym.spaces import Dict, Box, Discrete
from gym.vector import AsyncVectorEnv

class CustomEnv(gym.Env):
    observation_space = Box(low=0, high=255, shape=(84, 84), dtype=np.uint8)
    action_space = Dict({
        'fire': Discrete(2),
        'jump': Discrete(2),
        'move': Box(low=-1., high=1., shape=(2,), dtype=np.float32)
    })

    def reset(self):
        return self.observation_space.sample()

    def step(self, action):
        # Do something with action['fire'], action['jump'], action['move']
        observation = self.observation_space.sample()
        return observation, 0., False, {}

def make_env():
    return CustomEnv()

env = AsyncVectorEnv([make_env for _ in range(5)])
print(f'Action space: {env.action_space}')
env.reset()

actions = {
    'fire': np.array([1, 0, 0, 1, 1]),
    'jump': np.array([0, 0, 1, 1, 0]),
    'move': np.random.rand(5, 2)
}
observations, rewards, dones, infos = env.step(actions)
print(f'Observations shape: {observations.shape}')

# Action space: Dict(fire:MultiDiscrete([2 2 2 2 2]), jump:MultiDiscrete([2 2 2 2 2]), move:Box(-1.0, 1.0, (5, 2), float32))
# Observations shape: (5, 84, 84)

• Use batch_space instead of Tuple in VectorEnv
• Add the iterate utility function to iterate over items from a (batch) space
• Add tests
• Check if the action_space are all the same in all sub-environments

[lebrice]

Why not put the space as the first argument, and use a singledispatch callable here?
This would let users customize how this function should behave with their custom spaces, which isn't possible as it is.

[lebrice]

You're also essentially doing the same thing as a singledispatch here, but worse! :P

[tristandeleu]

It is definitely a prime candidate for singledispatch!
I was waiting for #2093, but I'll get the ball rolling and switch this function to singledispatch already.

[lebrice]

Nit: those docstrings and the doctest below still have the "items, space" ordering

[tristandeleu]

What is the status of this PR? Any timeline for this PR getting merged?

[tristandeleu]

Pinging @jkterry1

[jkterry1]

I'm just going to reply to all the notifications you sent in one place:

Everything regarding changes to the vector API is temporarily on hold until I get more time to develop a cohesive plan for it going forward. This is has been a lower priority for me than many other fixes to Gym at the moment. The documentation stuff isn't on hold per se, there's just been a long list of problems in getting the desired website fully operational and getting other parts of the text of the documentation written.

The status of all your other PRs is that I'm aware they exist and I need to go through them again myself for one reason or another and have not had the free time. Amongst other considerations, I'm working on 3 first author ICLR submissions. Gym wasn't maintained for years, things are going to take a little while to catch up.

[tristandeleu]

I see, perhaps you should get other maintainers on board to not be limited by the free time you can allocate to Gym, and so that you can offload some of those tasks.

For the particular case of this PR (to stay focused), these changes were discussed in #2279, and you were welcome to comment on those changes. You could rely on the community more instead of taking the responsibility to devise a plan by yourself if you are too busy.

[tristandeleu]

Any update on this PR? @jkterry1

[vwxyzjn]

The proposed changes make sense to me. I have a quick question.

print(env.action_space)
# Before: Tuple(Discrete(2), Discrete(2), Discrete(2), Discrete(2), Discrete(2))
# After: MultiDiscrete([2 2 2 2 2])

So what happened if you have Tuple(MultiDiscrete([2 2 2 2 2]), MultiDiscrete([2 2 2 2 2]))?

[tristandeleu]

You mean if the environment has action space Tuple(MultiDiscrete([2 2 2 2 2]), MultiDiscrete([2 2 2 2 2])) instead of Discrete(2)? Then it follows the rules of batch_space:

from gym.vector.utils.spaces import batch_space

space = Tuple((MultiDiscrete([2, 2, 2, 2, 2]), MultiDiscrete([2, 2, 2, 2, 2])))
batch_space(space, n=5)
# Tuple(Box([[0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]], [[1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]], (5, 5), int64), Box([[0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]], [[1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]], (5, 5), int64))

Just like for observation_space, any action_space can be batched if they are standard gym Space instances (i.e. Box, Discrete, Tuple, Dict, etc...).

[vwxyzjn]

No I meant if the envs.single_action_space is MultiDiscrete([2, 2, 2, 2, 2]), what is envs.action_space?

[tristandeleu]

Again, it follows the rules of batch_space

space = MultiDiscrete([2, 2, 2, 2, 2])
batch_space(space, n=5)
# Box([[0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]], [[1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]], (5, 5), int64)

[vwxyzjn]

I see. Thank you. The API makes sense. LGTM @jkterry1.

[tristandeleu]

A more complete example:

class CustomEnv(gym.Env):
    observation_space = Box(low=0., high=1., shape=(2,), dtype=np.float32)
    action_space = MultiDiscrete([2, 2, 2, 2, 2])

env = AsyncVectorEnv([lambda: CustomEnv() for _ in range(5)])
print(env.observation_space)
# Box([[0. 0.]
#  [0. 0.]
#  [0. 0.]
#  [0. 0.]
#  [0. 0.]], [[1. 1.]
#  [1. 1.]
#  [1. 1.]
#  [1. 1.]
#  [1. 1.]], (5, 2), float32)

print(env.action_space)
# Box([[0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]
#  [0 0 0 0 0]], [[1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]
#  [1 1 1 1 1]], (5, 5), int64)