Skip to content
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

Added list envs function (pretty print envs) #124

Merged
merged 3 commits into from
Nov 16, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion gymnasium/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -28,6 +28,7 @@
"spec",
"register",
"registry",
"pprint_registry",
# root files
"envs",
"spaces",
Expand Down
2 changes: 1 addition & 1 deletion gymnasium/envs/__init__.py
Original file line number Diff line number Diff line change
@@ -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()
Expand Down
56 changes: 56 additions & 0 deletions gymnasium/envs/registration.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import re
import sys
import warnings
from collections import defaultdict
from dataclasses import dataclass, field
from typing import (
Callable,
Expand Down Expand Up @@ -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
192 changes: 192 additions & 0 deletions tests/envs/test_pprint_registry.py
Original file line number Diff line number Diff line change
@@ -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"]