Spaces:
Configuration error
Configuration error
File size: 4,030 Bytes
82dd7b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | import re, random, importlib
from typing import Any, Union, List, Callable, Dict, Tuple, Optional
from dataclasses import dataclass, field
import textarena as ta
# Global environment registry
ENV_REGISTRY: Dict[str, Callable] = {}
@dataclass
class EnvSpec:
"""A specification for creating environments."""
id: str
entry_point: Callable
default_wrappers: Optional[List[ta.Wrapper]]
kwargs: Dict[str, Any] = field(default_factory=dict)
def make(self, **kwargs) -> Any:
"""Create an environment instance."""
all_kwargs = {**self.kwargs, **kwargs}
return self.entry_point(**all_kwargs)
def register(id: str, entry_point: Callable, default_wrappers: Optional[List[ta.Wrapper]]=None, **kwargs: Any):
"""Register an environment with a given ID."""
if id in ENV_REGISTRY:
raise ValueError(f"Environment {id} already registered.")
ENV_REGISTRY[id] = EnvSpec(id=id, entry_point=entry_point, default_wrappers=default_wrappers, kwargs=kwargs)
def register_with_versions(id: str, entry_point: Callable, wrappers: Optional[Dict[str, List[ta.Wrapper]]]=None, **kwargs: Any):
"""Register an environment with a given ID."""
if id in ENV_REGISTRY: raise ValueError(f"Environment {id} already registered.")
# first register default version
ENV_REGISTRY[id] = EnvSpec(id=id, entry_point=entry_point, default_wrappers=wrappers.get("default"), kwargs=kwargs)
for wrapper_version_key in list(wrappers.keys())+["-raw"]:
if wrapper_version_key=="default": continue
ENV_REGISTRY[f"{id}{wrapper_version_key}"] = EnvSpec(id=f"{id}{wrapper_version_key}", entry_point=entry_point, default_wrappers=wrappers.get(wrapper_version_key), kwargs=kwargs)
def pprint_registry_detailed():
"""Pretty print the registry with additional details like kwargs."""
if not ENV_REGISTRY:
print("No environments registered.")
else:
print("Detailed Registered Environments:")
for env_id, env_spec in ENV_REGISTRY.items():
print(f" - {env_id}:")
print(f" Entry Point: {env_spec.entry_point}")
print(f" Kwargs: {env_spec.kwargs}")
print(f" Wrappers: {env_spec.default_wrappers}")
def check_env_exists(env_id: str):
"""Check if an environment exists in the registry."""
if env_id not in ENV_REGISTRY:
raise ValueError(f"Environment {env_id} is not registered.")
else:
print(f"Environment {env_id} is registered.")
def make(env_id: Union[str, List[str]], **kwargs) -> Any:
"""Create an environment instance using the registered ID."""
# If env_id is a list, randomly select one environment ID
if isinstance(env_id, list):
if not env_id:
raise ValueError("Empty list of environment IDs provided.")
env_id = random.choice(env_id)
# Continue with the existing implementation
if env_id not in ENV_REGISTRY:
raise ValueError(f"Environment {env_id} not found in registry.")
env_spec = ENV_REGISTRY[env_id]
# Resolve the entry point if it's a string
if isinstance(env_spec.entry_point, str):
module_path, class_name = env_spec.entry_point.split(":")
try:
module = importlib.import_module(module_path)
env_class = getattr(module, class_name)
except (ModuleNotFoundError, AttributeError) as e:
raise ImportError(f"Could not import {module_path}.{class_name}. Error: {e}")
else:
env_class = env_spec.entry_point
env = env_class(**{**env_spec.kwargs, **kwargs})
# Dynamically attach the env_id
env.env_id = env_id
env.entry_point = env_spec.entry_point
# wrap the environment
if env_spec.default_wrappers is not None and len(env_spec.default_wrappers) > 0:
for wrapper in env_spec.default_wrappers:
env = wrapper(env)
return env
|