EnvPool's Async API

[vwxyzjn]

Description

This PR prototypes an async API like EnvPool, which looks like

import gymnasium
import time
import numpy as np
batch_size = 8
num_envs = 16

envs = gymnasium.vector.AsyncVectorEnv(
    [lambda: gymnasium.make("CartPole-v1") for _ in range(num_envs)],
    batch_size=batch_size
)
start_time = time.time()
envs.reset_async()
obs, rewards, terminateds, truncateds, infos = envs.recv()
print(obs.shape, infos)
envs.send(np.random.randint(envs.single_action_space.n, size=batch_size), infos["env_ids"])
obs, rewards, terminateds, truncateds, infos = envs.recv()
print(obs.shape, infos)
envs.send(np.random.randint(envs.single_action_space.n, size=batch_size), infos["env_ids"])
obs, rewards, terminateds, truncateds, infos = envs.recv()
print(obs.shape, infos)
envs.send(np.random.randint(envs.single_action_space.n, size=batch_size), infos["env_ids"])

(8, 4) {'env_ids': array([1, 3, 8, 2, 5, 4, 7, 0])}
(8, 4) {'env_ids': array([11, 10,  6,  9, 13, 12,  1, 14])}
(8, 4) {'env_ids': array([ 8,  2,  3,  5,  7,  4,  0, 15])}

Preliminary tests show that this API can provide further speed up. E.g., when using num_envs=64 and batch_size=16, this new async API offers a 2x speed improvement compared to regular Async mode. Further optimizations should be possible, too, but this PR mainly points out a proof-of-concept

https://github.com/vwxyzjn/Gymnasium/blob/9259d296d2097671a5f793bdb5164dd649c43b4a/a.py

Related to #93

Type of change

Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have run the pre-commit checks with pre-commit run --all-files (see CONTRIBUTING.md instructions to set it up)
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[pseudo-rnd-thoughts]

@vwxyzjn Could you change the PR to use the new-vector-api branch not main
Are we changing the API to use the envpool function calls or the envpool to use the gym function calls? I think it would be better to use the envpool. sail-sg/envpool#194

[RedTachyon]

So just for future travellers, to add some information from a discord conversation -- this is a proof-of-concept of an alternative implementation of AsyncVectorEnv which as I understand uses similar ideas to envpool, but doesn't use envpool itself as the library.

This wouldn't be a change on the API level, just the internals of the vectorization.

I think this isn't necessarily high priority (we don't need to have it merged by any concrete deadline), but if it's fleshed-out and profiled against the current implementation (also checked for identical outputs), I think it'd be a great addition.

[RedTachyon]

Closing to keep the PR backlog clean, the code will stay here if we want to refer back to it