Added list envs function (pretty print envs) (#124)

Co-authored-by: Mark Towers <[email protected]>

gymnasium/__init__.py

@@ -9,7 +9,7 @@
     RewardWrapper,
 )
 from gymnasium.spaces.space import Space
-from gymnasium.envs.registration import make, spec, register, registry
+from gymnasium.envs.registration import make, spec, register, registry, pprint_registry
 from gymnasium import envs, spaces, utils, vector, wrappers, error, logger
 
 import os
@@ -28,6 +28,7 @@
     "spec",
     "register",
     "registry",
+    "pprint_registry",
     # root files
     "envs",
     "spaces",

gymnasium/envs/__init__.py

@@ -1,5 +1,5 @@
 from gymnasium.envs.registration import load_env_plugins as _load_env_plugins
-from gymnasium.envs.registration import make, register, registry, spec
+from gymnasium.envs.registration import make, pprint_registry, register, registry, spec
 
 # Hook to load plugins from entry points
 _load_env_plugins()

gymnasium/envs/registration.py

@@ -6,6 +6,7 @@
 import re
 import sys
 import warnings
+from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import (
     Callable,
@@ -701,3 +702,58 @@ def spec(env_id: str) -> EnvSpec:
     else:
         assert isinstance(spec_, EnvSpec)
         return spec_
+
+
+def pprint_registry(
+    _registry: dict = registry,
+    max_rows: int = 10,
+    exclude_namespaces: Optional[List[str]] = None,
+) -> None:
+    """List the environments currently supported."""
+
+    # Defaultdict to store environment names according to namespace.
+    namespace_envs = defaultdict(lambda: [])
+    max_justify = float("-inf")
+    for env in _registry.values():
+        namespace, _, _ = parse_env_id(env.id)
+        print(namespace, env.id, env.entry_point)
+        if namespace is None:
+            # Since namespace is currently none, use regex to obtain namespace from entrypoints.
+            env_entry_point = re.sub(r":\w+", "", env.entry_point)
+            e_ep_split = env_entry_point.split(".")
+            if len(e_ep_split) >= 3:
+                # If namespace is of the format - gymnasium.envs.mujoco.ant_v4:AntEnv
+                # or gymnasium.envs.mujoco:HumanoidEnv
+                idx = 2
+                namespace = e_ep_split[idx]
+            elif len(e_ep_split) > 1:
+                # If namespace is of the format - shimmy.atari_env
+                idx = 1
+                namespace = e_ep_split[idx]
+            else:
+                # If namespace cannot be found, default to env id.
+                namespace = env.id
+        namespace_envs[namespace].append(env.id)
+        max_justify = max(max_justify, len(env.id))
+
+    # Iterate through each namespace and print environment alphabetically.
+    return_str = ""
+    for namespace, envs in namespace_envs.items():
+        # Ignore namespaces to exclude.
+        if exclude_namespaces is not None and namespace in exclude_namespaces:
+            continue
+        return_str += f"{'=' * 5} {namespace} {'=' * 5}\n"  # Print namespace.
+        num_columns = (
+            len(envs) // max_rows
+        ) + 1  # Calculate number of columns required.
+        # Reference: https://stackoverflow.com/a/33464001
+        for count, item in enumerate(sorted(envs), 1):
+            return_str += (
+                item.ljust(max_justify) + " "
+            )  # Print column with justification.
+            # Once all rows printed, switch to new column.
+            if count % num_columns == 0 or count == len(envs):
+                return_str += "\n"
+        return_str += "\n"
+
+    return return_str

tests/envs/test_pprint_registry.py

@@ -0,0 +1,192 @@
+import gymnasium as gym
+
+# To ignore the trailing whitespaces, will need flake to ignore this file.
+# flake8: noqa
+
+
+def test_pprint_custom_registry():
+    """Testing a registry different from default."""
+    a = {
+        "CartPole-v0": gym.envs.registry["CartPole-v0"],
+        "CartPole-v1": gym.envs.registry["CartPole-v1"],
+    }
+    out = gym.pprint_registry(a)
+
+    correct_out = """===== classic_control =====
+CartPole-v0 
+CartPole-v1 
+
+"""
+
+    assert out == correct_out
+
+
+def test_pprint_registry():
+    """Testing the default registry, with no changes."""
+    out = gym.pprint_registry()
+
+    correct_out = """===== classic_control =====
+Acrobot-v1                   
+CartPole-v0                  
+CartPole-v1                  
+MountainCar-v0               
+MountainCarContinuous-v0     
+Pendulum-v1                  
+
+===== box2d =====
+BipedalWalker-v3             
+BipedalWalkerHardcore-v3     
+CarRacing-v2                 
+LunarLander-v2               
+LunarLanderContinuous-v2     
+
+===== toy_text =====
+Blackjack-v1                 
+CliffWalking-v0              
+FrozenLake-v1                
+FrozenLake8x8-v1             
+Taxi-v3                      
+
+===== mujoco =====
+Ant-v2                       Ant-v3                       Ant-v4                       
+HalfCheetah-v2               HalfCheetah-v3               HalfCheetah-v4               
+Hopper-v2                    Hopper-v3                    Hopper-v4                    
+Humanoid-v2                  Humanoid-v3                  Humanoid-v4                  
+HumanoidStandup-v2           HumanoidStandup-v4           InvertedDoublePendulum-v2    
+InvertedDoublePendulum-v4    InvertedPendulum-v2          InvertedPendulum-v4          
+Pusher-v2                    Pusher-v4                    Reacher-v2                   
+Reacher-v4                   Swimmer-v2                   Swimmer-v3                   
+Swimmer-v4                   Walker2d-v2                  Walker2d-v3                  
+Walker2d-v4                  
+
+===== external =====
+GymV26Environment-v0         
+
+===== utils_envs =====
+RegisterDuringMakeEnv-v0     
+test.ArgumentEnv-v0          
+test.OrderlessArgumentEnv-v0 
+
+===== test =====
+test/NoHuman-v0              
+test/NoHumanNoRGB-v0         
+test/NoHumanOldAPI-v0        
+
+"""
+    assert out == correct_out
+
+
+def test_pprint_registry_exclude_namespaces():
+    """Testing the default registry, with no changes."""
+    out = gym.pprint_registry(
+        max_rows=20,
+        exclude_namespaces=["classic_control"],
+    )
+
+    correct_out = """===== box2d =====
+BipedalWalker-v3             
+BipedalWalkerHardcore-v3     
+CarRacing-v2                 
+LunarLander-v2               
+LunarLanderContinuous-v2     
+
+===== toy_text =====
+Blackjack-v1                 
+CliffWalking-v0              
+FrozenLake-v1                
+FrozenLake8x8-v1             
+Taxi-v3                      
+
+===== mujoco =====
+Ant-v2                       Ant-v3                       
+Ant-v4                       HalfCheetah-v2               
+HalfCheetah-v3               HalfCheetah-v4               
+Hopper-v2                    Hopper-v3                    
+Hopper-v4                    Humanoid-v2                  
+Humanoid-v3                  Humanoid-v4                  
+HumanoidStandup-v2           HumanoidStandup-v4           
+InvertedDoublePendulum-v2    InvertedDoublePendulum-v4    
+InvertedPendulum-v2          InvertedPendulum-v4          
+Pusher-v2                    Pusher-v4                    
+Reacher-v2                   Reacher-v4                   
+Swimmer-v2                   Swimmer-v3                   
+Swimmer-v4                   Walker2d-v2                  
+Walker2d-v3                  Walker2d-v4                  
+
+===== external =====
+GymV26Environment-v0         
+
+===== utils_envs =====
+RegisterDuringMakeEnv-v0     
+test.ArgumentEnv-v0          
+test.OrderlessArgumentEnv-v0 
+
+===== test =====
+test/NoHuman-v0              
+test/NoHumanNoRGB-v0         
+test/NoHumanOldAPI-v0        
+
+"""
+    assert out == correct_out
+
+
+def test_pprint_registry_no_entry_point():
+    """Test registry if there is environment with no entry point."""
+
+    gym.register("NoNamespaceEnv", "no-entry-point")
+    out = gym.pprint_registry()
+
+    correct_out = """===== classic_control =====
+Acrobot-v1                   
+CartPole-v0                  
+CartPole-v1                  
+MountainCar-v0               
+MountainCarContinuous-v0     
+Pendulum-v1                  
+
+===== box2d =====
+BipedalWalker-v3             
+BipedalWalkerHardcore-v3     
+CarRacing-v2                 
+LunarLander-v2               
+LunarLanderContinuous-v2     
+
+===== toy_text =====
+Blackjack-v1                 
+CliffWalking-v0              
+FrozenLake-v1                
+FrozenLake8x8-v1             
+Taxi-v3                      
+
+===== mujoco =====
+Ant-v2                       Ant-v3                       Ant-v4                       
+HalfCheetah-v2               HalfCheetah-v3               HalfCheetah-v4               
+Hopper-v2                    Hopper-v3                    Hopper-v4                    
+Humanoid-v2                  Humanoid-v3                  Humanoid-v4                  
+HumanoidStandup-v2           HumanoidStandup-v4           InvertedDoublePendulum-v2    
+InvertedDoublePendulum-v4    InvertedPendulum-v2          InvertedPendulum-v4          
+Pusher-v2                    Pusher-v4                    Reacher-v2                   
+Reacher-v4                   Swimmer-v2                   Swimmer-v3                   
+Swimmer-v4                   Walker2d-v2                  Walker2d-v3                  
+Walker2d-v4                  
+
+===== external =====
+GymV26Environment-v0         
+
+===== utils_envs =====
+RegisterDuringMakeEnv-v0     
+test.ArgumentEnv-v0          
+test.OrderlessArgumentEnv-v0 
+
+===== test =====
+test/NoHuman-v0              
+test/NoHumanNoRGB-v0         
+test/NoHumanOldAPI-v0        
+
+===== NoNamespaceEnv =====
+NoNamespaceEnv               
+
+"""
+    assert out == correct_out
+
+    del gym.envs.registry["NoNamespaceEnv"]