Skip to content

Adds native vector implementation #599

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Conversation

pseudo-rnd-thoughts
Copy link
Member

@pseudo-rnd-thoughts pseudo-rnd-thoughts commented Mar 25, 2025

Projects like SampleFactory and EnvPool have provided C++ based vectorisators for ALE, significantly increasing the number of frames per second that can be sampled.

However, EnvPool is no longer maintained and to make it easier and compatible with Gymnasium's vector API, this PR adds a EnvPool based (though highly highly simplified) vector environment.

The implementation has four main parts

  1. preprocessed_env.hpp - Runs the standard preprocessing computation for Atari (frame stack, frame skip, maxpooling, resizing, etc)
  2. async_vectorizers.hpp - Takes a list of preprocessed environments that are run asynchronously that we push actions to then return a collection of results once computed.
  3. ale_vector_python_interface.cpp/hpp - PyBind11 bindings to convert the vectorizer's list of environment data into a batched version that is easier to use by users
  4. vector_env.py - This is a python class for easily using the vectorizer.

Known Problems

  • For MacOS ARM (not x86), the ALE vector and AtariPreprocess are NOT equivalent for all timesteps, however, the difference in only 1 or 2 pixels with a difference of 1. The issue is almost certainly with OpenCV with a different implement between MacOS ARM and x86.

Features not included in the PR (todo later)

  • This PR only supports next-step autoreset mode, future PRs could look to add disabled autoreset mode (relatively easy) and same-step autoreset mode (more difficult)
  • XLA support - The EnvPool implementation adds support for jit accelerating the environment, however, the jax foreign function interface has changed significantly since the EnvPool therefore will need to be implemented from scratch.
  • Testing - This PR is minimally tested, particularly checking that agents can be successfully trained

@pseudo-rnd-thoughts pseudo-rnd-thoughts changed the title Simple vector Adds native vector implementation Apr 1, 2025
@jjshoots
Copy link
Member

Reviewing it soon, thanks for the work on this!

Copy link
Member

@jjshoots jjshoots left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Admittedly, a lot of it flew over my head and I never ran the code. But thus doesn't seem to break any existing functionality so I thibk it's good to go. Any downstream bugs can be fixed as they come a la open source.

@pseudo-rnd-thoughts
Copy link
Member Author

@jjshoots Thanks for the review, it was complex code and took me a long time to implement so don't worry if you don't fully understand

@pseudo-rnd-thoughts pseudo-rnd-thoughts merged commit c800316 into Farama-Foundation:master Apr 26, 2025
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants