textpolicy 0.0.1__py3-none-any.whl → 0.1.1__py3-none-any.whl

textpolicy/cli.py

@@ -0,0 +1,67 @@
+"""
+Minimal CLI entry point for TextPolicy.
+
+Design goals:
+- Keep it tiny and dependency-free (argparse only)
+- Provide a single high-signal command for students: `validate`
+- Exit non-zero on failure for CI integration
+
+This CLI is intentionally small; a config-driven runner can be added later.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from typing import Any, Dict
+
+from .validate import validate_installation
+
+
+def _cmd_validate(args: argparse.Namespace) -> int:
+    """Run installation validation and print results.
+
+    Uses the programmatic validate_installation() so behavior stays consistent.
+    """
+    report: Dict[str, Any] = validate_installation(verbose=not args.json)
+    if args.json:
+        print(json.dumps(report, indent=2))
+    # Exit code communicates status for CI/automation
+    return 0 if report.get("status") == "ok" else 1
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Create the top-level CLI parser.
+
+    We avoid subcommand bloat; a single `validate` command covers health checks.
+    Default behavior is `validate` when no subcommand is provided for convenience.
+    """
+    parser = argparse.ArgumentParser(prog="textpolicy", description="TextPolicy command-line tools")
+    subparsers = parser.add_subparsers(dest="command")
+
+    p_validate = subparsers.add_parser("validate", help="Validate installation and environment")
+    p_validate.add_argument("--json", action="store_true", help="Output machine-readable JSON report")
+    p_validate.set_defaults(func=_cmd_validate)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entrypoint. Defaults to `validate` if no command is provided."""
+    parser = build_parser()
+    # If no args given, behave like `textpolicy validate` for a quick health check
+    if argv is None:
+        argv = sys.argv[1:]
+    if not argv:
+        argv = ["validate"]
+    args = parser.parse_args(argv)
+    if not hasattr(args, "func"):
+        parser.print_help()
+        return 2
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+

textpolicy/environment/__init__.py

@@ -0,0 +1,79 @@
+# mlx_rl/environment/__init__.py
+"""
+Modular environment implementation for MLX-RL.
+"""
+
+from .environment import (
+    # Base classes
+    Environment,
+    EnvironmentAdapter,
+    
+    # Adapters  
+    GymAdapter,
+    
+    # Factory functions
+    create_environment,
+    register_environment,
+    list_registered_environments,
+    is_gymnasium_available,
+    
+    # Registry
+    ENVIRONMENT_REGISTRY
+)
+
+from .vectorized import (
+    # Vectorized environment classes
+    VectorizedEnvironment,
+    VectorizedCollector,
+    
+    # Factory function
+    make_vectorized_env,
+)
+
+# Task suite registry for text generation environments
+from .task_suites import (
+    register_task_suite,
+    list_task_suites,
+    get_task_suite,
+)
+# Re-export text generation environments and helpers for public API access
+from .text_generation import (
+    TextGenerationEnvironment,
+    TextGenerationEnv,
+    create_text_generation_test_env,
+    validate_learning_progress,
+)
+
+__all__ = [
+    # Base classes
+    "Environment", 
+    "EnvironmentAdapter",
+    
+    # Adapters
+    "GymAdapter",
+    
+    # Vectorized environment classes
+    "VectorizedEnvironment",
+    "VectorizedCollector",
+    
+    # Factory functions
+    "create_environment",
+    "register_environment",
+    "list_registered_environments", 
+    "is_gymnasium_available",
+    "make_vectorized_env",
+    
+    # Task suite registry
+    "register_task_suite",
+    "list_task_suites",
+    "get_task_suite",
+    
+    # Registry
+    "ENVIRONMENT_REGISTRY",
+    
+    # Text generation environments and helpers
+    "TextGenerationEnvironment",
+    "TextGenerationEnv",
+    "create_text_generation_test_env",
+    "validate_learning_progress",
+] 

textpolicy/environment/base.py

@@ -0,0 +1,110 @@
+# mlx_rl/environment/base.py
+"""
+Base environment interface and protocols for MLX-RL.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Tuple, Protocol
+
+
+class Environment(ABC):
+    """
+    Unified environment interface for all environment types.
+    
+    This abstract base class defines the contract that all environments
+    must implement to work with MLX-RL agents and trainers.
+    """
+
+    @abstractmethod
+    def reset(self) -> Tuple[Any, Dict[str, Any]]:
+        """
+        Reset environment and return initial observation.
+        
+        Returns:
+            Tuple of (observation, info) following gymnasium API
+        """
+        pass
+
+    @abstractmethod
+    def step(self, action: Any) -> Dict[str, Any]:
+        """
+        Take action and return step result.
+        
+        Args:
+            action: Action to take in the environment
+            
+        Returns:
+            Dict with keys: observation, reward, terminated, truncated, info
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def observation_space(self) -> Any:
+        """
+        Observation space specification.
+        
+        Returns:
+            Space object describing valid observations
+        """
+        pass
+
+    @property
+    @abstractmethod
+    def action_space(self) -> Any:
+        """
+        Action space specification.
+        
+        Returns:
+            Space object describing valid actions
+        """
+        pass
+
+    def clone(self) -> 'Environment':
+        """
+        Create a clone of this environment for multiprocessing.
+        
+        Default implementation raises NotImplementedError.
+        Subclasses should override if they support cloning.
+        
+        Returns:
+            New instance of the same environment
+        """
+        # Default clone implementation: use deepcopy to support multiprocessing clones
+        import copy
+        return copy.deepcopy(self)
+
+    def render(self, mode: str = "human") -> Any:
+        """
+        Render the environment.
+        
+        Args:
+            mode: Rendering mode (e.g., "human", "rgb_array")
+            
+        Returns:
+            Rendered output (depends on mode)
+        """
+        # Default implementation does nothing
+        pass
+
+    def close(self):
+        """Clean up environment resources."""
+        # Default implementation does nothing
+        pass
+
+
+class EnvironmentAdapter(Protocol):
+    """
+    Protocol for environment adapters.
+    
+    Adapters convert external environment APIs (gym, dm_env, etc.)
+    to the unified Environment interface.
+    """
+    
+    def __init__(self, env_spec: Any, **kwargs):
+        """Initialize adapter with environment specification."""
+        ...
+    
+    def clone(self) -> Environment:
+        """Create a clone for multiprocessing."""
+        ... 

textpolicy/environment/environment.py

@@ -0,0 +1,46 @@
+# mlx_rl/environment/environment.py
+"""
+Unified environment module that coordinates all environment components.
+
+Environment system for MLX-RL.
+
+This module provides a unified interface for different environment types
+(gymnasium, dm_env, custom environments) with adapter patterns for compatibility.
+
+Design principles:
+1. Unified Environment interface for all environment types
+2. Adapter pattern for external environment libraries
+3. Factory system for easy environment creation and registration
+4. Support for multiprocessing through cloning
+5. Extensible design for adding new environment types
+"""
+
+# Import all components for unified API
+from .base import Environment, EnvironmentAdapter
+from .gym import GymAdapter
+from .factory import (
+    create_environment,
+    register_environment,
+    list_registered_environments,
+    is_gymnasium_available,
+    ENVIRONMENT_REGISTRY
+)
+
+# Re-export everything for compatibility
+__all__ = [
+    # Base classes
+    "Environment",
+    "EnvironmentAdapter",
+    
+    # Adapters
+    "GymAdapter",
+    
+    # Factory functions
+    "create_environment",
+    "register_environment", 
+    "list_registered_environments",
+    "is_gymnasium_available",
+    
+    # Registry
+    "ENVIRONMENT_REGISTRY"
+] 

textpolicy/environment/factory.py

@@ -0,0 +1,103 @@
+# mlx_rl/environment/factory.py
+"""
+Environment factory and registration system.
+"""
+
+from typing import Dict, Callable, Union
+from .base import Environment
+from .gym import GymAdapter
+
+# Environment registry for dynamic loading
+ENVIRONMENT_REGISTRY: Dict[str, Callable] = {}
+
+
+def register_environment(name: str, factory_func: Callable):
+    """
+    Register an environment factory function.
+    
+    Args:
+        name: Environment name (will be lowercased)
+        factory_func: Function that returns Environment instance
+    """
+    ENVIRONMENT_REGISTRY[name.lower()] = factory_func
+
+
+def create_environment(env_spec: Union[str, Environment], **kwargs) -> Environment:
+    """
+    Create environment from specification.
+    
+    Args:
+        env_spec: Either environment name string or Environment instance
+        **kwargs: Additional arguments for environment creation
+        
+    Returns:
+        Environment instance
+        
+    Examples:
+        # Create from string (uses GymAdapter)
+        env = create_environment("CartPole-v1")
+        
+        # Create with custom parameters
+        env = create_environment("LunarLander-v2", continuous=True)
+        
+        # Pass through existing environment
+        env = create_environment(my_custom_env)
+    """
+    if isinstance(env_spec, str):
+        # String specification - look up in registry first
+        env_name = env_spec.lower()
+        
+        if env_name in ENVIRONMENT_REGISTRY:
+            # Use registered factory
+            return ENVIRONMENT_REGISTRY[env_name](**kwargs)
+        else:
+            # Default to gymnasium adapter
+            return GymAdapter(env_spec, **kwargs)
+    
+    elif isinstance(env_spec, Environment):
+        # Already an environment instance
+        if kwargs:
+            raise ValueError("Cannot pass kwargs when env_spec is already an Environment instance")
+        return env_spec
+    
+    else:
+        raise TypeError(f"env_spec must be str or Environment, got {type(env_spec)}")
+
+
+def list_registered_environments() -> list[str]:
+    """
+    List all registered environment names.
+    
+    Returns:
+        List of registered environment names
+    """
+    return list(ENVIRONMENT_REGISTRY.keys())
+
+
+def is_gymnasium_available() -> bool:
+    """
+    Check if gymnasium is available for import.
+    
+    Returns:
+        True if gymnasium can be imported, False otherwise
+    """
+    try:
+        return True
+    except ImportError:
+        return False
+
+
+# Register some common environments by default
+def _register_defaults():
+    """Register default environment factories."""
+    
+    # Gymnasium environments (if available)
+    if is_gymnasium_available():
+        register_environment("cartpole", lambda **kwargs: GymAdapter("CartPole-v1", **kwargs))
+        register_environment("cartpole-v1", lambda **kwargs: GymAdapter("CartPole-v1", **kwargs))
+        register_environment("lunarlander", lambda **kwargs: GymAdapter("LunarLander-v2", **kwargs))
+        register_environment("lunarlander-v2", lambda **kwargs: GymAdapter("LunarLander-v2", **kwargs))
+
+
+# Register defaults on import
+_register_defaults() 

textpolicy/environment/gym.py

@@ -0,0 +1,106 @@
+# mlx_rl/environment/gym.py
+"""
+Gymnasium environment adapter for MLX-RL.
+"""
+
+from typing import Any, Dict, Tuple
+import gymnasium as gym
+from .base import Environment
+
+
+class GymAdapter(Environment):
+    """
+    Adapter for gymnasium environments.
+    
+    Converts gymnasium environments to the unified Environment interface.
+    Handles both old and new gymnasium APIs for maximum compatibility.
+    """
+
+    def __init__(self, env_name: str, **kwargs):
+        """
+        Initialize gymnasium environment.
+        
+        Args:
+            env_name: Name of the gymnasium environment (e.g., "CartPole-v1")
+            **kwargs: Additional arguments passed to gym.make()
+        """
+        self.env_name = env_name
+        self.env_kwargs = kwargs
+        self.env = gym.make(env_name, **kwargs)
+
+    def clone(self) -> 'GymAdapter':
+        """
+        Create a new instance of the same environment.
+        
+        This is essential for multiprocessing where each worker
+        needs its own environment instance.
+        
+        Returns:
+            New GymAdapter instance with same configuration
+        """
+        return GymAdapter(self.env_name, **self.env_kwargs)
+
+    def reset(self) -> Tuple[Any, Dict[str, Any]]:
+        """
+        Reset environment and return initial observation.
+        
+        Handles both old gymnasium API (returns tuple) and
+        newer API (returns observation, info).
+        
+        Returns:
+            Tuple of (observation, info)
+        """
+        result = self.env.reset()
+        if isinstance(result, tuple):
+            return result  # New API: (obs, info)
+        else:
+            return result, {}  # Old API: just obs
+
+    def step(self, action: Any) -> Dict[str, Any]:
+        """
+        Take action and return step result in unified format.
+        
+        Args:
+            action: Action to take (format depends on action space)
+            
+        Returns:
+            Dictionary with observation, reward, terminated, truncated, info
+        """
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        return {
+            "observation": obs,
+            "reward": reward,
+            "terminated": terminated,
+            "truncated": truncated,
+            "info": info
+        }
+
+    @property
+    def observation_space(self):
+        """Get observation space from underlying gymnasium environment."""
+        return self.env.observation_space
+
+    @property
+    def action_space(self):
+        """Get action space from underlying gymnasium environment."""
+        return self.env.action_space
+
+    def render(self, mode: str = "human") -> Any:
+        """
+        Render the environment.
+        
+        Args:
+            mode: Rendering mode ("human", "rgb_array", etc.)
+            
+        Returns:
+            Rendered output (depends on mode and environment)
+        """
+        return self.env.render()
+
+    def close(self):
+        """Close the underlying gymnasium environment."""
+        self.env.close()
+
+    def __repr__(self) -> str:
+        """String representation for debugging."""
+        return f"GymAdapter(env_name='{self.env_name}', kwargs={self.env_kwargs})" 

textpolicy/environment/task_suites.py

@@ -0,0 +1,51 @@
+"""
+Task suite registry for TextGenerationEnvironment.
+
+This module provides a minimal registry that maps suite names to loader
+functions returning lists of TextGenerationTask instances. It enables
+customizable evaluation suites without hardcoding them in the environment.
+
+Notes:
+- Keep this registry lightweight and dependency-free.
+- Default suites are registered from text_generation.py to avoid import cycles.
+- A file-backed loader (JSON/YAML) can be added later if needed.
+"""
+
+from __future__ import annotations
+
+from typing import Callable, Dict, List, Any, Optional
+
+
+# Registry type: name -> callable that returns a List[TextGenerationTask]
+_TASK_SUITE_REGISTRY: Dict[str, Callable[[], List[Any]]] = {}
+
+
+def register_task_suite(name: str, loader: Callable[[], List[Any]]) -> None:
+    """
+    Register a task suite by name.
+
+    Args:
+        name: Suite identifier (e.g., "basic", "challenging").
+        loader: Callable that returns a list of TextGenerationTask instances.
+    """
+    if not callable(loader):
+        raise TypeError("loader must be callable and return a list of tasks")
+    _TASK_SUITE_REGISTRY[name] = loader
+
+
+def get_task_suite(name: str) -> Optional[List[Any]]:
+    """
+    Load a registered task suite by name.
+
+    Returns None if the suite is not registered.
+    """
+    loader = _TASK_SUITE_REGISTRY.get(name)
+    if loader is None:
+        return None
+    return loader()
+
+
+def list_task_suites() -> List[str]:
+    """List available task suite names."""
+    return sorted(_TASK_SUITE_REGISTRY.keys())
+