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

textpolicy/rollout/rollout.py

@@ -0,0 +1,142 @@
+# textpolicy/rollout/rollout.py
+"""
+Main rollout coordinator and public interface.
+"""
+
+from typing import List, Callable, Any
+from .worker import RolloutWorker
+from .runner import RolloutRunner
+from .aggregator import BufferAggregator
+from .strategy import create_strategy
+from textpolicy.buffer import Buffer
+
+
+class RolloutCoordinator:
+    """
+    Coordinator for rollout collection.
+    
+    Provides unified interface for:
+    - Single-process rollouts (for debugging/small-scale)
+    - Multi-process rollouts (for production/performance)
+    - Strategy management and worker coordination
+    """
+    
+    def __init__(
+        self,
+        env_fn: Callable[[], Any],
+        policy_fn: Callable[[], Any],
+        algorithm: str,
+        num_workers: int = 0,
+        max_steps: int = 1000,
+        max_episodes: int = 100
+    ):
+        """
+        Initialize rollout coordinator.
+        
+        Args:
+            env_fn: Function that creates environment instances
+            policy_fn: Function that creates policy instances
+            algorithm: Algorithm name ('ppo', 'grpo', etc.)
+            num_workers: Number of worker processes (0 = single-process mode)
+            max_steps: Maximum steps per rollout
+            max_episodes: Maximum episodes to buffer
+        """
+        self.env_fn = env_fn
+        self.policy_fn = policy_fn
+        self.algorithm = algorithm
+        self.num_workers = num_workers
+        self.max_steps = max_steps
+        
+        # Create strategy for algorithm
+        self.strategy = create_strategy(algorithm)
+        
+        # Setup for multi-process or single-process mode
+        if num_workers > 0:
+            self._setup_multiprocess(max_episodes)
+        else:
+            self._setup_singleprocess()
+    
+    def _setup_multiprocess(self, max_episodes: int):
+        """Setup multi-process rollout collection."""
+        self.aggregator = BufferAggregator(self.num_workers, max_episodes)
+        self.workers: List[RolloutWorker] = []
+        
+        # Create and start worker processes
+        for i in range(self.num_workers):
+            worker = RolloutWorker(
+                env_fn=self.env_fn,
+                policy_fn=self.policy_fn,
+                strategy=self.strategy,
+                max_steps=self.max_steps
+            )
+            self.workers.append(worker)
+            self.aggregator.add_worker(worker, i)
+            worker.start()
+    
+    def _setup_singleprocess(self):
+        """Setup single-process rollout collection."""
+        self.aggregator = None
+        self.workers = []
+        
+        # Create single runner for direct use
+        env = self.env_fn()
+        policy = self.policy_fn()
+        self.runner = RolloutRunner(env, policy, self.strategy, self.max_steps)
+    
+    def collect(self) -> Buffer:
+        """
+        Collect rollout data.
+        
+        Returns:
+            Buffer containing collected episodes
+        """
+        if self.num_workers > 0:
+            return self._collect_multiprocess()
+        else:
+            return self._collect_singleprocess()
+    
+    def _collect_multiprocess(self) -> Buffer:
+        """Collect data using multiple worker processes."""
+        # Request rollouts from all workers
+        for worker in self.workers:
+            worker.collect_async()
+        
+        # Wait for and consume results
+        while not self.aggregator.ready(min_episodes=self.num_workers):
+            self.aggregator.consume_all()
+        
+        # Return aggregated buffer
+        # Note: In practice, trainer would manage this more carefully
+        return self.aggregator.buffer
+    
+    def _collect_singleprocess(self) -> Buffer:
+        """Collect data using single process."""
+        return self.runner.collect()
+    
+    def close(self):
+        """Cleanup resources."""
+        if self.workers:
+            for worker in self.workers:
+                worker.close()
+
+
+# Public API functions for external use
+def create_rollout_coordinator(
+    env_fn: Callable[[], Any],
+    policy_fn: Callable[[], Any], 
+    algorithm: str,
+    **kwargs
+) -> RolloutCoordinator:
+    """
+    Factory function for creating rollout coordinators.
+    
+    Args:
+        env_fn: Environment factory function
+        policy_fn: Policy factory function
+        algorithm: Algorithm name
+        **kwargs: Additional configuration options
+        
+    Returns:
+        RolloutCoordinator instance
+    """
+    return RolloutCoordinator(env_fn, policy_fn, algorithm, **kwargs)

textpolicy/rollout/runner.py

@@ -0,0 +1,280 @@
+# textpolicy/rollout/runner.py
+"""
+Core rollout collection engine.
+"""
+
+from typing import Callable, Dict, Optional, Tuple
+import mlx.core as mx # type: ignore
+from textpolicy.buffer import Buffer
+from .base import RolloutStrategy, DEFAULT_MAX_STEPS
+
+
+class RolloutRunner:
+    """
+    Collects rollouts from a single environment using a policy.
+    
+    Designed for MLX and Apple Silicon:
+    - Minimized Python overhead in the collection loop
+    - MLX array conversions performed once per step
+    - Python lists for storage to reduce overhead
+    - Policy inference on GPU/ANE when available
+    """
+    
+    def __init__(
+        self,
+        env,
+        policy: Optional[Callable[[mx.array], Tuple[mx.array, Dict[str, mx.array]]]] = None,
+        strategy: Optional[RolloutStrategy] = None,
+        max_steps: int = DEFAULT_MAX_STEPS,
+        agent = None  # Alternative API: pass agent instead of policy/strategy
+    ) -> None:
+        """
+        Initialize rollout runner.
+        
+        Args:
+            env: Environment instance (must implement gym interface)
+            policy: Policy function that takes obs and returns (action, extras)
+            strategy: RolloutStrategy defining algorithm-specific behavior
+            max_steps: Maximum steps per rollout collection
+            agent: Alternative API - Agent object containing policy and rollout_strategy
+        """
+        self.env = env
+        
+        # Support both direct policy/strategy and agent-based initialization
+        if agent is not None:
+            # Extract policy and strategy from agent for backward compatibility with tests
+            self.policy = agent.policy
+            self.strategy = agent.rollout_strategy
+        else:
+            # Direct policy/strategy initialization (current API)
+            if policy is None or strategy is None:
+                raise ValueError("Must provide either agent or both policy and strategy")
+            self.policy = policy
+            self.strategy = strategy
+            
+        self.max_steps = max_steps
+        self.buffer = Buffer(max_episodes=10)
+        self.step_count = 0  # Track total steps collected for test compatibility
+
+    def _normalize_step_result(self, step_result):
+        """
+        Normalize Environment.step results to a tuple
+        (next_obs, reward, terminated, truncated, info).
+
+        Enforces dict-shaped step results per Environment contract.
+        Raises TypeError for tuple-based results.
+        """
+        if not isinstance(step_result, dict):
+            raise TypeError(
+                "Environment.step must return a dict with keys: observation, reward, "
+                "terminated, truncated, info. Tuple returns are not supported."
+            )
+        return (
+            step_result.get("observation"),
+            step_result.get("reward"),
+            step_result.get("terminated"),
+            step_result.get("truncated"),
+            step_result.get("info", {}),
+        )
+
+    def collect(self) -> Buffer:
+        """
+        Run one complete rollout collection.
+        
+        MLX efficiency guidelines:
+        - Policy runs on GPU/ANE via MLX arrays
+        - Buffer stores as Python lists
+        - Batch array conversions reduce memory transfers
+        - Strategy handles algorithm differences
+        
+        Returns:
+            Buffer containing collected episodes
+        """
+        obs, _ = self.env.reset()
+
+        # Batch observations for efficient MLX conversion
+        # Collect multiple observations before converting to MLX arrays
+        batch_size = min(32, self.max_steps)  # Batch size for array conversion
+        obs_batch = []
+
+        for step in range(self.max_steps):
+            obs_batch.append(obs)
+            
+            # Process batch when full or on final step
+            if len(obs_batch) == batch_size or step == self.max_steps - 1:
+                # Single MLX array conversion for the entire batch
+                try:
+                    import numpy as np
+                    obs_batch_np = np.array(obs_batch)
+                    obs_batch_mx = mx.array(obs_batch_np)
+                except (ValueError, TypeError):
+                    # Fallback to individual conversion if batch conversion fails
+                    obs_batch_mx = mx.stack([mx.array(o) for o in obs_batch])
+                
+                # Process each observation in the batch
+                for i, obs_single in enumerate(obs_batch):
+                    # Extract single observation from batch
+                    obs_mx = obs_batch_mx[i] if len(obs_batch) > 1 else obs_batch_mx
+                    
+                    # Policy forward pass - runs on GPU/ANE
+                    action_mx, extra = self.strategy.select_action(self.policy, obs_mx)
+
+                    # Convert action back to Python format (once per step)
+                    if action_mx.ndim == 0:
+                        action = action_mx.item()  # Scalar action (discrete environments)
+                    else:
+                        action = action_mx.tolist()  # Vector action (continuous environments)
+
+                    # Environment step (CPU-bound). Normalize to a standard tuple.
+                    step_result = self.env.step(action)
+                    next_obs, reward, done, trunc, info = self._normalize_step_result(step_result)
+
+                    # Store transition using strategy-specific logic
+                    # Strategy handles filtering and algorithm-specific data
+                    self.strategy.store_transition(
+                        buffer=self.buffer,
+                        obs=obs_single,  # Use original observation
+                        act=action,
+                        rew=reward,
+                        next_obs=next_obs,
+                        done=done,
+                        timeout=trunc,
+                        **extra  # Algorithm-specific data (logprob, value, etc.)
+                    )
+
+                    # Handle episode boundaries
+                    if done or trunc:
+                        obs, _ = self.env.reset()
+                    else:
+                        obs = next_obs
+                    
+                    # Break if we've reached max steps
+                    if step >= self.max_steps - 1:
+                        break
+                
+                # Clear batch for next iteration
+                obs_batch = []
+
+        return self.buffer
+    
+    def collect_episode(self, deterministic=False):
+        """
+        Collect a single episode and return as trajectory (backward compatibility).
+        
+        Args:
+            deterministic: Whether to use deterministic policy actions
+        
+        Returns:
+            List of transition dictionaries for test compatibility
+        """
+        # Reset buffer for clean episode collection
+        self.buffer.clear()
+        
+        # Handle both old-style (obs only) and new-style (obs, info) reset
+        reset_result = self.env.reset()
+        if isinstance(reset_result, tuple):
+            obs, _ = reset_result
+        else:
+            obs = reset_result
+        trajectory = []
+        
+        for step in range(self.max_steps):
+            # Convert observation to MLX array
+            obs_mx = mx.array(obs)
+            
+            # Policy forward pass - pass through deterministic flag to policy
+            if hasattr(self.policy, '__call__'):
+                # If policy supports deterministic parameter, use it
+                try:
+                    action_mx, extra = self.policy(obs_mx, deterministic=deterministic)
+                except TypeError:
+                    # Fallback: use strategy select_action which handles deterministic behavior
+                    action_mx, extra = self.strategy.select_action(self.policy, obs_mx)
+            else:
+                action_mx, extra = self.strategy.select_action(self.policy, obs_mx)
+            
+            # Convert action back to Python format
+            if action_mx.ndim == 0:
+                action = action_mx.item()  # Scalar action (discrete environments)
+            else:
+                action = action_mx.tolist()  # Vector action (continuous environments)
+            
+            # Environment step: normalize dict/tuple to a standard tuple
+            step_result = self.env.step(action)
+            next_obs, reward, done, trunc, info = self._normalize_step_result(step_result)
+            
+            # Build transition dictionary for test compatibility
+            transition = {
+                'obs': obs,
+                'act': action,
+                'rew': reward,
+                'next_obs': next_obs,
+                'done': done or trunc,  # Combine terminated/truncated for backward compatibility
+                **extra  # Include logprob, value, entropy from strategy
+            }
+            trajectory.append(transition)
+            
+            # Track step count for tests
+            self.step_count += 1
+            
+            # Handle episode boundaries
+            if done or trunc:
+                break
+            else:
+                obs = next_obs
+                
+        return trajectory
+    
+    def collect_rollout(self, num_episodes: int, collect_stats=False):
+        """
+        Collect multiple episodes (backward compatibility).
+        
+        Args:
+            num_episodes: Number of episodes to collect
+            collect_stats: Whether to collect statistics (for test compatibility)
+            
+        Returns:
+            List of episode trajectories
+        """
+        trajectories = []
+        for _ in range(num_episodes):
+            trajectory = self.collect_episode()
+            trajectories.append(trajectory)
+        
+        # Store trajectories for statistics if requested
+        if collect_stats:
+            self._last_trajectories = trajectories
+            
+        return trajectories
+    
+    def get_statistics(self):
+        """
+        Get rollout statistics (for test compatibility).
+        
+        Returns:
+            Dictionary of rollout statistics
+        """
+        if not hasattr(self, '_last_trajectories'):
+            return {
+                'total_episodes': 0,
+                'total_steps': self.step_count,
+                'avg_episode_length': 0,
+                'avg_episode_reward': 0
+            }
+        
+        total_episodes = len(self._last_trajectories)
+        total_steps = sum(len(traj) for traj in self._last_trajectories)
+        avg_episode_length = total_steps / max(total_episodes, 1)
+        
+        total_reward = sum(
+            sum(transition['rew'] for transition in traj) 
+            for traj in self._last_trajectories
+        )
+        avg_episode_reward = total_reward / max(total_episodes, 1)
+        
+        return {
+            'total_episodes': total_episodes,
+            'total_steps': total_steps,
+            'avg_episode_length': avg_episode_length,
+            'avg_episode_reward': avg_episode_reward
+        } 

textpolicy/rollout/strategy.py

@@ -0,0 +1,208 @@
+# textpolicy/rollout/strategy.py
+"""
+Algorithm-specific rollout strategies.
+"""
+
+from typing import Callable, Dict, Any, Tuple
+import mlx.core as mx # type: ignore
+from textpolicy.buffer import Buffer
+from .base import RolloutStrategy, validate_transition_data
+
+
+class PPOStrategy(RolloutStrategy):
+    """
+    Rollout strategy for Proximal Policy Optimization (PPO).
+    
+    PPO requires:
+    - Action probabilities (logprob) for policy gradient
+    - Value function estimates for advantage calculation
+    - Policy and value function trained together
+    
+    Expected policy output:
+        (action, {"logprob": mx.array, "value": mx.array, "entropy": mx.array})
+    
+    Stored in buffer: obs, act, rew, next_obs, done, timeout, logprob, value
+    Filtered out: entropy (computed during training, not stored)
+    """
+    
+    def select_action(self, policy: Callable, obs: mx.array) -> Tuple[mx.array, Dict[str, Any]]:
+        """
+        Select action using PPO policy.
+        
+        PPO TRAINING BEHAVIOR: Uses stochastic policy (deterministic=False) for exploration
+        during training data collection. This provides the variety of experiences needed
+        for robust policy learning. Evaluation uses deterministic=True for consistent
+        performance measurement.
+        
+        Args:
+            policy: Policy function returning (action, extras)
+            obs: MLX array observation
+            
+        Returns:
+            action: Selected action as MLX array (sampled stochastically)
+            extras: Dict with logprob, value, entropy
+        """
+        # Use stochastic policy for training data collection (explore during training, evaluate deterministically)
+        return policy(obs, deterministic=False)
+
+    def store_transition(self, buffer: Buffer, **data) -> None:
+        """
+        Store PPO transition data in buffer.
+        
+        Filters data to include only what the buffer supports and PPO needs.
+        Validates required fields are present.
+        
+        Args:
+            buffer: Buffer instance to store data
+            **data: Transition data including obs, act, rew, etc.
+        """
+        # Validate and filter transition data
+        filtered_data = validate_transition_data(data)
+        
+        # Remove entropy if present (not stored, computed during training)
+        filtered_data.pop('entropy', None)
+        
+        # Store in buffer
+        buffer.add(**filtered_data) # type: ignore
+
+
+class GRPOStrategy(RolloutStrategy):
+    """
+    Rollout strategy for Group Relative Policy Optimization (GRPO).
+    
+    GRPO characteristics:
+    - No value function required (uses group-relative advantages)
+    - Only needs action probabilities for policy gradient
+    - Advantage computed relative to group performance
+    
+    Expected policy output:
+        (action, {"logprob": mx.array, "entropy": mx.array})
+    
+    Stored in buffer: obs, act, rew, next_obs, done, timeout, logprob
+    Filtered out: value (not used), entropy (computed during training)
+    """
+    
+    def select_action(self, policy: Callable, obs: mx.array) -> Tuple[mx.array, Dict[str, Any]]:
+        """
+        Select action using GRPO policy.
+        
+        GRPO TRAINING BEHAVIOR: Uses stochastic policy (deterministic=False) for exploration
+        during training data collection, consistent with PPO approach.
+        
+        Args:
+            policy: Policy function returning (action, extras)
+            obs: MLX array observation
+            
+        Returns:
+            action: Selected action as MLX array (sampled stochastically)
+            extras: Dict with logprob, entropy (no value function)
+        """
+        # Use stochastic policy for training data collection
+        return policy(obs, deterministic=False)
+
+    def store_transition(self, buffer: Buffer, **data) -> None:
+        """
+        Store GRPO transition data in buffer.
+        
+        GRPO doesn't use value functions, so value data is filtered out.
+        Only stores what's needed for group-relative advantage computation.
+        
+        Args:
+            buffer: Buffer instance to store data
+            **data: Transition data including obs, act, rew, etc.
+        """
+        # Validate and filter transition data
+        filtered_data = validate_transition_data(data)
+        
+        # Remove fields not used by GRPO
+        filtered_data.pop('value', None)    # No value function in GRPO
+        filtered_data.pop('entropy', None)  # Computed during training
+        
+        # Store in buffer
+        buffer.add(**filtered_data) # type: ignore
+
+
+class DreamerV3Strategy(RolloutStrategy):
+    """
+    Rollout strategy for DreamerV3.
+    
+    DreamerV3 is model-based RL that requires:
+    - Real environment interactions for world model learning
+    - Action probabilities for actor training
+    - State representations for RSSM dynamics learning
+    
+    Expected policy output:
+        (action, {"logprob": mx.array, "value": mx.array, "entropy": mx.array, 
+                  "state": mx.array, "embed": mx.array})
+    
+    Stored in buffer: obs, act, rew, next_obs, done, timeout, logprob, value, state, embed
+    Filtered out: entropy (computed during training)
+    """
+    
+    def select_action(self, policy: Callable, obs: mx.array) -> Tuple[mx.array, Dict[str, Any]]:
+        """
+        Select action using DreamerV3 policy.
+        
+        DreamerV3 uses the world model's posterior state for action selection.
+        During training, we use stochastic policy for exploration.
+        
+        Args:
+            policy: DreamerV3 policy function returning (action, extras)
+            obs: MLX array observation
+            
+        Returns:
+            action: Selected action as MLX array (sampled stochastically)
+            extras: Dict with logprob, value, entropy, state, embed
+        """
+        # Use stochastic policy during training for exploration
+        return policy(obs, deterministic=False)
+
+    def store_transition(self, buffer: Buffer, **data) -> None:
+        """
+        Store DreamerV3 transition data in buffer.
+        
+        DreamerV3 stores additional state representations and embeddings
+        needed for world model learning and imagination.
+        
+        Args:
+            buffer: Buffer instance to store data
+            **data: Transition data including obs, act, rew, state, embed, etc.
+        """
+        # Validate and filter transition data
+        filtered_data = validate_transition_data(data)
+        
+        # Remove entropy if present (not stored, computed during training)
+        filtered_data.pop('entropy', None)
+        
+        # Store in buffer (including DreamerV3-specific fields like state, embed)
+        buffer.add(**filtered_data) # type: ignore
+
+
+# Strategy registry for factory pattern
+STRATEGY_REGISTRY = {
+    'ppo': PPOStrategy,
+    'grpo': GRPOStrategy,
+    'gspo': GRPOStrategy,  # Add alias for backwards compatibility
+    'dreamerv3': DreamerV3Strategy,
+}
+
+
+def create_strategy(algorithm: str) -> RolloutStrategy:
+    """
+    Factory function for creating rollout strategies.
+    
+    Args:
+        algorithm: Algorithm name ('ppo', 'grpo', etc.)
+        
+    Returns:
+        RolloutStrategy instance
+        
+    Raises:
+        ValueError: If algorithm is not supported
+    """
+    if algorithm not in STRATEGY_REGISTRY:
+        available = list(STRATEGY_REGISTRY.keys())
+        raise ValueError(f"Unknown algorithm '{algorithm}'. Available: {available}")
+    
+    strategy_class = STRATEGY_REGISTRY[algorithm]
+    return strategy_class()