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

textpolicy/buffer/sampling.py

@@ -0,0 +1,438 @@
+# textpolicy/buffer/sampling.py
+"""
+Buffer sampling methods for training data retrieval.
+
+Designed for MLX and Apple Silicon with efficient tensor conversions.
+"""
+
+from typing import List, Dict
+import random
+import mlx.core as mx # type: ignore
+from .episode import Episode
+
+
+class BufferSampler:
+    """
+    Handles all data sampling operations for the buffer.
+    
+    Provides multiple sampling strategies:
+    - Full buffer sampling
+    - Latest N steps sampling
+    - Episode-based sampling
+    
+    Uses MLX tensor operations and Apple Silicon-friendly patterns.
+    """
+    
+    def __init__(self, episodes: List[Episode]):
+        """
+        Initialize sampler with episode storage reference.
+        
+        Args:
+            episodes: Reference to list of complete episodes
+        """
+        self.episodes = episodes
+    
+    def sample_all(self) -> Dict[str, mx.array]:
+        """
+        Sample all stored episodes as a single concatenated batch.
+        
+        Returns all transitions in chronological order:
+        - Oldest episode → Newest episode  
+        - Each episode: first step → last step
+        
+        Returns:
+            Dict of MLX arrays with all transitions
+            
+        Raises:
+            ValueError: If buffer is empty
+        """
+        if not self.episodes:
+            raise ValueError("Buffer empty. No episodes to sample.")
+        
+        # Concatenate all episodes by collecting data directly (bypass Episode validation)
+        # Episodes can have different optional fields, so collect what exists across all episodes
+        all_obs = []
+        all_act = []
+        all_rew = []
+        all_next_obs = []
+        all_done = []
+        all_timeout = []
+        all_logprob = []
+        all_value = []
+        all_entropy = []
+        
+        # Only include optional fields that exist in ALL episodes for consistent sampling
+        # This matches the buffer's "all-or-nothing" design philosophy per episode
+        has_logprob = all(episode.logprob is not None for episode in self.episodes)
+        has_value = all(episode.value is not None for episode in self.episodes)
+        has_entropy = all(episode.entropy is not None for episode in self.episodes)
+        
+        # Collect transitions from all episodes
+        for episode in self.episodes:
+            for i in range(len(episode)):
+                all_obs.append(episode.obs[i])
+                all_act.append(episode.act[i])
+                all_rew.append(episode.rew[i])
+                all_next_obs.append(episode.next_obs[i])
+                all_done.append(episode.done[i])
+                all_timeout.append(episode.timeout[i])
+                
+                # Only collect optional fields that exist in ALL episodes
+                if has_logprob:
+                    all_logprob.append(episode.logprob[i])
+                if has_value:
+                    all_value.append(episode.value[i])
+                if has_entropy:
+                    all_entropy.append(episode.entropy[i])
+        
+        # Create Episode directly with collected data (bypassing validation during construction)
+        all_transitions = Episode()
+        all_transitions.obs = all_obs
+        all_transitions.act = all_act
+        all_transitions.rew = all_rew
+        all_transitions.next_obs = all_next_obs
+        all_transitions.done = all_done
+        all_transitions.timeout = all_timeout
+        
+        # Only set optional fields if they exist in any episode
+        if has_logprob:
+            all_transitions.logprob = all_logprob
+        if has_value:
+            all_transitions.value = all_value
+        if has_entropy:
+            all_transitions.entropy = all_entropy
+        
+        return all_transitions.to_tensor_dict()
+    
+    def sample_latest_steps(self, n: int) -> Dict[str, mx.array]:
+        """
+        Sample the N most recent transitions across episodes.
+        
+        Useful for on-policy RL algorithms that use recent experience for stable training.
+        
+        Args:
+            n: Number of steps to sample (must be > 0)
+            
+        Returns:
+            Dict of MLX arrays with the latest n steps in chronological order
+            (oldest → newest, ensuring temporal consistency)
+            
+        Raises:
+            ValueError: If buffer is empty, n <= 0, or not enough steps
+        """
+        if not self.episodes:
+            raise ValueError("Buffer is empty. No recent episodes to sample.")
+        if n <= 0:
+            raise ValueError("Number of steps (n) to sample must be greater than 0.")
+        
+        # Collect steps in reverse chronological order (newest first)
+        steps = []
+        for episode in reversed(self.episodes):
+            for i in reversed(range(len(episode))):
+                step_dict = {
+                    "obs": episode.obs[i],
+                    "act": episode.act[i],
+                    "rew": episode.rew[i], 
+                    "next_obs": episode.next_obs[i],
+                    "done": episode.done[i],
+                    "timeout": episode.timeout[i],
+                }
+                
+                # Add optional fields if present
+                if episode.logprob is not None:
+                    step_dict["logprob"] = episode.logprob[i]
+                if episode.value is not None:
+                    step_dict["value"] = episode.value[i]
+                if episode.entropy is not None:
+                    step_dict["entropy"] = episode.entropy[i]
+                
+                steps.append(step_dict)
+                
+                # Stop when we have enough steps
+                if len(steps) >= n:
+                    break
+            
+            if len(steps) >= n:
+                break
+        
+        if not steps:
+            raise ValueError("No steps available to sample")
+        
+        # Reverse to get chronological order (oldest → newest)
+        chronological_steps = list(reversed(steps[:n]))
+        
+        # Batch convert to MLX arrays
+        # Collect values first, then convert in a single pass for memory efficiency
+        batch = {}
+        for key in chronological_steps[0].keys():
+            # Extract all values for this key at once
+            values = [step[key] for step in chronological_steps]
+            # Single batch conversion is more efficient than individual mx.array() calls
+            batch[key] = mx.stack([mx.array(v) for v in values])
+        
+        return batch
+    
+    def sample_episodes(self, k: int, order: str = 'asc') -> Dict[str, mx.array]:
+        """
+        Sample up to k complete episodes.
+        
+        Useful for episode-based training or evaluation analysis.
+        
+        Args:
+            k: Number of episodes to sample (must be > 0)
+            order: 'asc' for oldest first, 'desc' for newest first
+            
+        Returns:
+            Dict of MLX arrays with concatenated transitions from selected episodes
+            
+        Raises:
+            ValueError: If buffer is empty, k <= 0, or invalid order
+        """
+        if not self.episodes:
+            raise ValueError("Buffer is empty. No episodes to sample.")
+        if k <= 0:
+            raise ValueError("k must be positive.")
+        if order not in ('asc', 'desc'):
+            raise ValueError("order must be 'asc' or 'desc'")
+        
+        # Select episodes based on order
+        if order == 'asc':
+            selected_episodes = self.episodes[:k]  # Oldest k episodes
+        else:  # 'desc'
+            selected_episodes = self.episodes[-k:]  # Newest k episodes
+        
+        if not selected_episodes:
+            raise ValueError("No episodes matched the criteria.")
+        
+        # Concatenate all steps from selected episodes
+        all_transitions = Episode()
+        for episode in selected_episodes:
+            for i in range(len(episode)):
+                all_transitions.append(
+                    obs=episode.obs[i],
+                    act=episode.act[i],
+                    rew=episode.rew[i],
+                    next_obs=episode.next_obs[i], 
+                    done=episode.done[i],
+                    timeout=episode.timeout[i],
+                    logprob=episode.logprob[i] if episode.logprob is not None else None,
+                    value=episode.value[i] if episode.value is not None else None,
+                    entropy=episode.entropy[i] if episode.entropy is not None else None
+                )
+        
+        return all_transitions.to_tensor_dict()
+
+    def sample_sequences(
+        self,
+        batch_size: int,
+        seq_len: int,
+        recent_first: bool = True,
+        drop_incomplete: bool = True,
+        dreamerv3_mode: bool = False,
+    ) -> Dict[str, mx.array]:
+        """
+        Sample contiguous sequences of length `seq_len`.
+
+        Standard mode: Sample without crossing episode boundaries (for most algorithms).
+        DreamerV3 mode: Sample across episode boundaries to include terminals (for continue head supervision).
+
+        Args:
+            batch_size: Number of sequences to sample.
+            seq_len: Length of each sequence (T). Must be > 0.
+            recent_first: Prefer sampling from most recent episodes first.
+            drop_incomplete: If True, skip episodes shorter than seq_len (standard mode only).
+            dreamerv3_mode: If True, sample across episode boundaries to include terminals.
+
+        Returns:
+            Dict of MLX arrays with keys: obs, act, rew, next_obs, done, timeout
+            and optional keys (logprob, value, entropy) included only if present.
+
+        Raises:
+            ValueError: If buffer is empty or inputs are invalid or not enough data.
+        """
+        if not self.episodes:
+            raise ValueError("Buffer is empty. No episodes to sample.")
+        if batch_size <= 0 or seq_len <= 0:
+            raise ValueError("batch_size and seq_len must be positive.")
+
+        if dreamerv3_mode:
+            return self._sample_dreamerv3_sequences(batch_size, seq_len, recent_first)
+
+        # Choose episode order per recency preference.
+        episodes_iter: List[Episode]
+        episodes_iter = list(reversed(self.episodes)) if recent_first else list(self.episodes)
+
+        # Collect one latest contiguous window per episode until batch is filled.
+        sequences = []
+        episodes_used: List[Episode] = []
+
+        for ep in episodes_iter:
+            n = len(ep)
+            if n < seq_len:
+                if drop_incomplete:
+                    continue
+                # Padding/masking path intentionally not implemented for simplicity/perf.
+                continue
+
+            # DreamerV3 expects sequences sampled from throughout episodes, not only tails.
+            # Avoid tail-bias that would overrepresent terminal steps. Sample a random
+            # start index within the episode. This preserves recency at the episode
+            # level via episodes_iter and reduces bias within each episode.
+            max_start = n - seq_len
+            start = random.randint(0, max_start) if max_start > 0 else 0
+            end = start + seq_len  # ensure fixed-length window [start, start+seq_len)
+            seq = {
+                'obs': ep.obs[start:end],
+                'act': ep.act[start:end],
+                'rew': ep.rew[start:end],
+                'next_obs': ep.next_obs[start:end],
+                'done': ep.done[start:end],
+                'timeout': ep.timeout[start:end],
+            }
+
+            # Optional fields: only include if present in the episode
+            if ep.logprob is not None:
+                seq['logprob'] = ep.logprob[start:end]
+            if ep.value is not None:
+                seq['value'] = ep.value[start:end]
+            if ep.entropy is not None:
+                seq['entropy'] = ep.entropy[start:end]
+
+            sequences.append(seq)
+            episodes_used.append(ep)
+            if len(sequences) >= batch_size:
+                break
+
+        if not sequences:
+            raise ValueError("No sequences available to sample.")
+
+        # Determine optional keys that exist across all sampled sequences for consistent batching.
+        all_keys = set(sequences[0].keys())
+        for s in sequences[1:]:
+            all_keys &= set(s.keys())
+
+        # Convert to MLX arrays with shape [B, T, ...]
+        batch: Dict[str, mx.array] = {}
+        for key in all_keys:
+            # First convert each sequence to [T, ...]
+            per_seq = []
+            for s in sequences:
+                # Efficient batch conversion: one mx.array() per time step then stack along time
+                # to minimize Python overhead while keeping explicit control of dimensions.
+                per_seq.append(mx.stack([mx.array(v) for v in s[key]], axis=0))  # [T, ...]
+            # Stack sequences along batch dimension
+            batch[key] = mx.stack(per_seq, axis=0)  # [B, T, ...]
+
+        return batch
+
+    def _sample_dreamerv3_sequences(
+        self, 
+        batch_size: int, 
+        seq_len: int, 
+        recent_first: bool
+    ) -> Dict[str, mx.array]:
+        """
+        Sample sequences for DreamerV3 by concatenating episodes to form continuous trajectories.
+        
+        Unlike standard sampling, this allows sequences to span episode boundaries,
+        ensuring episode terminals appear within training sequences for continue head supervision.
+        
+        This mirrors DreamerV3's original replay buffer behavior where episodes are stored
+        continuously and sampling naturally includes episode boundaries.
+        """
+        import random
+        
+        # Create continuous trajectory by concatenating recent episodes
+        episodes_iter = list(reversed(self.episodes)) if recent_first else list(self.episodes)
+        
+        # Concatenate episodes into a continuous stream
+        continuous_data = {
+            'obs': [], 'act': [], 'rew': [], 'next_obs': [], 'done': [], 'timeout': [],
+            'logprob': [], 'value': [], 'entropy': []
+        }
+        
+        for ep in episodes_iter:
+            continuous_data['obs'].extend(ep.obs)
+            continuous_data['act'].extend(ep.act) 
+            continuous_data['rew'].extend(ep.rew)
+            continuous_data['next_obs'].extend(ep.next_obs)
+            continuous_data['done'].extend(ep.done)
+            continuous_data['timeout'].extend(ep.timeout)
+            
+            # Optional fields - extend with zeros if episode doesn't have them
+            if ep.logprob is not None:
+                continuous_data['logprob'].extend(ep.logprob)
+            else:
+                continuous_data['logprob'].extend([0.0] * len(ep.obs))
+                
+            if ep.value is not None:
+                continuous_data['value'].extend(ep.value)
+            else:
+                continuous_data['value'].extend([0.0] * len(ep.obs))
+                
+            if ep.entropy is not None:
+                continuous_data['entropy'].extend(ep.entropy)
+            else:
+                continuous_data['entropy'].extend([0.0] * len(ep.obs))
+        
+        total_steps = len(continuous_data['obs'])
+        if total_steps < seq_len:
+            raise ValueError(f"Not enough continuous data: {total_steps} steps < {seq_len} required")
+        
+        # Sample random sequences from the continuous trajectory
+        sequences = []
+        for _ in range(batch_size):
+            # Random start position that allows full sequence
+            max_start = total_steps - seq_len
+            start = random.randint(0, max_start) if max_start > 0 else 0
+            end = start + seq_len
+            
+            # Extract sequence
+            seq = {}
+            for key in ['obs', 'act', 'rew', 'next_obs', 'done', 'timeout', 'logprob', 'value', 'entropy']:
+                seq[key] = continuous_data[key][start:end]
+            
+            sequences.append(seq)
+        
+        # Convert to batch format [batch_size, seq_len, ...] using the same efficient conversion as standard mode
+        batch = {}
+        for key in sequences[0].keys():
+            per_seq = []
+            for s in sequences:
+                # Convert sequence to mx.array and stack along time dimension
+                per_seq.append(mx.stack([mx.array(v) for v in s[key]], axis=0))  # [T, ...]
+            # Stack sequences along batch dimension  
+            batch[key] = mx.stack(per_seq, axis=0)  # [B, T, ...]
+
+        return batch
+    
+    def get_episode_statistics(self) -> Dict[str, float]:
+        """
+        Get statistical information about stored episodes.
+        
+        Returns:
+            Dictionary with episode statistics for analysis
+        """
+        if not self.episodes:
+            return {
+                'episode_count': 0,
+                'total_steps': 0,
+                'mean_episode_length': 0.0,
+                'min_episode_length': 0,
+                'max_episode_length': 0,
+                'total_reward': 0.0,
+                'mean_episode_reward': 0.0
+            }
+        
+        episode_lengths = [len(ep) for ep in self.episodes]
+        episode_rewards = [sum(ep.rew) for ep in self.episodes]
+        
+        return {
+            'episode_count': len(self.episodes),
+            'total_steps': sum(episode_lengths),
+            'mean_episode_length': sum(episode_lengths) / len(episode_lengths),
+            'min_episode_length': min(episode_lengths),
+            'max_episode_length': max(episode_lengths),
+            'total_reward': sum(episode_rewards),
+            'mean_episode_reward': sum(episode_rewards) / len(episode_rewards)
+        } 

textpolicy/buffer/storage.py

@@ -0,0 +1,255 @@
+# textpolicy/buffer/storage.py
+"""
+Buffer storage and capacity management.
+
+Handles episode storage, capacity limits, and episode lifecycle management.
+Designed for efficient memory usage on Apple Silicon.
+"""
+
+from typing import List, Dict, Any
+from .episode import Episode
+
+
+class BufferStorage:
+    """
+    Manages episode storage with capacity limits and lifecycle.
+    
+    Features:
+    - FIFO episode eviction when capacity exceeded
+    - Episode validation before storage
+    - Efficient storage for multiprocessing scenarios
+    - Memory-conscious design for Apple Silicon
+    """
+    
+    def __init__(self, max_episodes: int = 100):
+        """
+        Initialize buffer storage.
+        
+        Args:
+            max_episodes: Maximum number of complete episodes to store.
+                         Oldest episodes are dropped when capacity is exceeded.
+        """
+        self.max_episodes = max_episodes
+        self.episodes: List[Episode] = []
+        self.current_episode = Episode()
+    
+    def add_transition(
+        self,
+        obs: Any,
+        act: Any, 
+        rew: Any,
+        next_obs: Any,
+        done: bool,
+        timeout: bool = False,
+        **kwargs  # Additional fields like logprob, value, entropy
+    ):
+        """
+        Add a transition to the current episode.
+        
+        Completes and stores the episode if done or timeout is True.
+        
+        Args:
+            obs: Observation
+            act: Action taken
+            rew: Reward received
+            next_obs: Next observation
+            done: Boolean indicating episode termination
+            timeout: Boolean indicating truncation (e.g. time limit)
+            **kwargs: Optional fields (logprob, value, entropy)
+        """
+        # Add transition to current episode
+        self.current_episode.append(
+            obs=obs, act=act, rew=rew, next_obs=next_obs,
+            done=done, timeout=timeout, **kwargs
+        )
+        
+        # Complete episode if terminated
+        if done or timeout:
+            self._complete_current_episode()
+    
+    def _complete_current_episode(self):
+        """
+        Complete the current episode and start a new one.
+        
+        Validates episode before storage and enforces capacity limits.
+        """
+        # Validate episode before storage
+        if len(self.current_episode) > 0:
+            self.current_episode.validate_consistency()
+            
+            # Debug: episode data quality
+            self._debug_episode_data(self.current_episode)
+            
+            # Add to storage
+            self.episodes.append(self.current_episode)
+            
+            # Enforce capacity limit (FIFO eviction)
+            if len(self.episodes) > self.max_episodes:
+                self.episodes.pop(0)  # Remove oldest episode
+        
+        # Start new episode
+        self.current_episode = Episode()
+    
+    def _debug_episode_data(self, episode):
+        """Debug what episodes look like - only show every 10th episode."""
+        if not hasattr(self, 'episode_debug_count'):
+            self.episode_debug_count = 0
+        
+        self.episode_debug_count += 1
+        
+        # Only debug every 50th episode to avoid spam
+        if self.episode_debug_count % 50 == 1:
+            try:
+                # Type checking and safe conversion
+                episode_count = int(self.episode_debug_count)
+                
+                rewards = episode.rew if hasattr(episode, 'rew') else []
+                # Convert all rewards to float before summing to handle mixed types
+                try:
+                    numeric_rewards = [float(r) for r in rewards]
+                    total_reward = sum(numeric_rewards)
+                except (TypeError, ValueError):
+                    # Fallback: filter out non-numeric rewards
+                    numeric_rewards = []
+                    for r in rewards:
+                        try:
+                            numeric_rewards.append(float(r))
+                        except (TypeError, ValueError):
+                            continue
+                    total_reward = sum(numeric_rewards) if numeric_rewards else 0.0
+                episode_length = len(episode)
+                
+                print(f"\nEPISODE DEBUG (Episode #{episode_count}):")
+                print(f"  Episode length: {episode_length}")
+                print(f"  Total reward: {total_reward:.2f}")
+                
+                if rewards:
+                    # Ensure rewards are numeric before formatting
+                    try:
+                        first_rewards = [float(r) for r in rewards[:10]]
+                        last_reward = float(rewards[-1])
+                        print(f"  Reward sequence: {first_rewards}...")
+                        print(f"  Last reward: {last_reward:.2f}")
+                    except (TypeError, ValueError) as reward_error:
+                        print(f"  Reward formatting error: {reward_error}")
+                        print(f"  Raw rewards: {rewards[:10]}...")
+                
+                # Check termination
+                if hasattr(episode, 'done') and episode.done:
+                    final_done = episode.done[-1] if episode.done else False
+                    print(f"  Final done: {final_done}")
+                
+                if hasattr(episode, 'timeout') and episode.timeout:
+                    final_timeout = episode.timeout[-1] if episode.timeout else False
+                    print(f"  Final timeout: {final_timeout}")
+                    
+            except Exception as e:
+                print(f"  Episode debug error: {e}")
+                # More detailed error info
+                import traceback
+                print(f"  Error details: {traceback.format_exc()}")
+    
+    def add_episode_from_dict(self, data: Dict[str, Any]):
+        """
+        Reconstruct and add an episode from serialized dictionary.
+        
+        Used for multiprocessing: worker serializes episode to dict,
+        trainer deserializes and adds to buffer.
+        
+        Args:
+            data: Dictionary containing episode data from episode.to_dict()
+                  Must include: obs, act, rew, next_obs, done, timeout
+                  Optional: logprob, value, entropy
+        
+        Raises:
+            ValueError: If episode data is invalid or inconsistent
+        """
+        # Create new episode from dictionary
+        episode = Episode()
+        
+        # Set required fields
+        episode.obs = data['obs']
+        episode.act = data['act'] 
+        episode.rew = data['rew']
+        episode.next_obs = data['next_obs']
+        episode.done = data['done']
+        episode.timeout = data['timeout']
+        
+        # Set optional fields if present
+        episode.logprob = data.get('logprob', None)
+        episode.value = data.get('value', None) 
+        episode.entropy = data.get('entropy', None)
+        
+        # Validate consistency before adding
+        episode.validate_consistency()
+        
+        # Add to storage
+        self.episodes.append(episode)
+        
+        # Enforce capacity limit
+        if len(self.episodes) > self.max_episodes:
+            self.episodes.pop(0)
+    
+    def clear(self):
+        """
+        Clear all stored episodes and reset current episode.
+        
+        Used to reset buffer state between training runs.
+        """
+        self.episodes.clear()
+        self.current_episode = Episode()
+    
+    def ready(self, min_episodes: int = 1) -> bool:
+        """
+        Check if buffer has enough complete episodes for training.
+        
+        Args:
+            min_episodes: Minimum number of episodes required
+            
+        Returns:
+            True if buffer has at least min_episodes complete episodes
+        """
+        return len(self.episodes) >= min_episodes
+    
+    def total_steps(self) -> int:
+        """
+        Calculate total number of steps across all episodes.
+        
+        Returns:
+            Sum of steps in all complete episodes
+        """
+        return sum(len(episode) for episode in self.episodes)
+    
+    def get_episodes(self) -> List[Episode]:
+        """
+        Get read-only access to stored episodes.
+        
+        Returns:
+            List of complete episodes (does not include current incomplete episode)
+        """
+        return self.episodes.copy()
+    
+    def __len__(self) -> int:
+        """Return total number of steps in buffer."""
+        return self.total_steps()
+    
+    @property
+    def episode_count(self) -> int:
+        """Number of complete episodes currently stored."""
+        return len(self.episodes)
+    
+    def get_storage_info(self) -> Dict[str, Any]:
+        """
+        Get detailed storage information for debugging.
+        
+        Returns:
+            Dictionary with storage statistics
+        """
+        return {
+            'episode_count': len(self.episodes),
+            'max_episodes': self.max_episodes,
+            'total_steps': self.total_steps(),
+            'current_episode_steps': len(self.current_episode),
+            'episode_lengths': [len(ep) for ep in self.episodes],
+            'capacity_usage': len(self.episodes) / self.max_episodes if self.max_episodes > 0 else 0.0
+        }