stable-baselines3 icon indicating copy to clipboard operation
stable-baselines3 copied to clipboard

WIP: Support for Gymnasium Graph endpoint for sb3 #2121 :

Open dhruvmalik007 opened this issue 5 months ago • 4 comments

Abstract change

Resolves the issue of supporting the gymnasium.graph instances not being support in sb3 .

Description

This pull request introduces support for Graph observation spaces in the Stable-Baselines3 environment checker, along with minor code cleanups and new test cases. The most significant changes include adding compatibility for GraphInstance observations, extending validation logic, and introducing new tests to ensure robustness.

Support for Graph Observation Spaces:

  • Added imports for Graph and GraphInstance from gymnasium.spaces.graph in stable_baselines3/common/env_checker.py to handle graph-based observation spaces.
  • Introduced _is_graph_space utility function to check if an observation space is a Graph. This helps in conditionally handling graph-related logic.
  • Updated _check_obs to validate GraphInstance observations, ensuring compatibility with single-value tuples and direct instances. [1] [2]
  • Modified _check_returned_values to handle Graph spaces during reset and step validations, ensuring backward compatibility with legacy tuple formats. [1] [2]
  • Enhanced check_env to skip NaN checks for Graph spaces, as they are not applicable.

Testing Enhancements:

  • Added a new ProperGraphEnv class in tests/test_env_checker.py to simulate environments with valid and invalid Graph observation spaces.
  • Introduced test cases to verify check_env functionality with Graph spaces, including valid scenarios and commented-out cases for invalid configurations.

Minor Code Cleanups:

  • Removed redundant blank lines and comments across multiple files for better readability. [1] [2] [3]
  • Updated obs_as_tensor in stable_baselines3/common/utils.py to handle unrecognized observation types gracefully by returning them as-is.

These changes ensure that Stable-Baselines3 can handle environments with graph-based observation spaces while maintaining backward compatibility and improving code quality.

  • [X] I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

  • [X] 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 change)
  • [ ] Documentation (update in the documentation)

Checklist

  • [X] I've read the CONTRIBUTION guide (required)
  • [x] I have updated the changelog accordingly (required).
  • [ ] My change requires a change to the documentation.
  • [X] I have updated the tests accordingly (required for a bug fix or a new feature).
  • [ ] I have updated the documentation accordingly.
  • [ ] I have opened an associated PR on the SB3-Contrib repository (if necessary)
  • [ ] I have opened an associated PR on the RL-Zoo3 repository (if necessary)
  • [X] I have reformatted the code using make format (required)
  • [X] I have checked the codestyle using make check-codestyle and make lint (required)
  • [X] I have ensured make pytest and make type both pass. (required)
  • [X] I have checked that the documentation builds using make doc (required) This pull request introduces support for Graph observation spaces in Stable-Baselines3, enabling the use of graph-based environments. The changes include updates to handle Graph spaces throughout the codebase, the addition of new utility functions, and the introduction of a base graph neural network class and a simple implementation for feature extraction. Below are the most important changes grouped by theme:

dhruvmalik007 avatar May 14 '25 14:05 dhruvmalik007