phantomrt 0.1.0__py3-none-any.whl

atlas/core/dynamics.py

@@ -0,0 +1,217 @@
+"""
+Dynamics Model: The Physics Engine
+
+This is the HEART of the world model.
+It learns HOW THE WORLD EVOLVES over time.
+
+Instead of discrete layers, we use a Neural ODE —
+a continuous differential equation that describes
+how the state changes:
+
+    dx/dt = f(x, action)
+
+This is fundamentally different from transformers:
+- Continuous, not discrete
+- Learns dynamics, not patterns
+- Can simulate arbitrary time horizons
+- Naturally handles variable-speed events
+"""
+
+import torch
+import torch.nn as nn
+from torchdiffeq import odeint
+
+
+class DynamicsFunction(nn.Module):
+    """
+    Neural network that defines the dynamics: dx/dt = f(x, action)
+    
+    Input:  (state, action) concatenated
+    Output: derivative dx/dt (rate of change of state)
+    
+    The network learns the PHYSICS of the environment —
+    how objects move, interact, and change over time.
+    """
+
+    def __init__(self, state_dim: int, action_dim: int, hidden_dim: int = 512, num_layers: int = 3):
+        super().__init__()
+        
+        self.state_dim = state_dim
+        self.action_dim = action_dim
+        
+        # Build dynamics network
+        layers = []
+        input_dim = state_dim + action_dim
+        
+        for i in range(num_layers):
+            output_dim = hidden_dim if i < num_layers - 1 else state_dim
+            layers.extend([
+                nn.Linear(input_dim if i == 0 else hidden_dim, hidden_dim),
+                nn.LayerNorm(hidden_dim),
+                nn.SiLU(),
+            ])
+            # Skip connection every 2 layers for gradient flow
+            if i > 0 and i % 2 == 0:
+                layers.append(SkipConnection(hidden_dim))
+        
+        # Final layer outputs the derivative
+        layers.append(nn.Linear(hidden_dim, state_dim))
+        
+        self.net = nn.Sequential(*layers)
+        
+        # Initialize last layer with small weights
+        # (prevents explosive dynamics at the start)
+        last_layer = self.net[-1]
+        nn.init.xavier_uniform_(last_layer.weight, gain=0.01)
+        nn.init.zeros_(last_layer.bias)
+    
+    def forward(self, state: torch.Tensor, t: torch.Tensor, action: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the derivative dx/dt.
+        
+        Args:
+            state: [batch, state_dim] current state
+            t: scalar or [batch] current time (needed for ODE solver)
+            action: [batch, action_dim] action being taken
+            
+        Returns:
+            dx_dt: [batch, state_dim] rate of change of state
+        """
+        # Concatenate state and action
+        x = torch.cat([state, action], dim=-1)
+        
+        # Compute derivative
+        dx_dt = self.net(x)
+        
+        # Optional: clamp derivatives to prevent explosion
+        dx_dt = torch.clamp(dx_dt, min=-10.0, max=10.0)
+        
+        return dx_dt
+
+
+class SkipConnection(nn.Module):
+    """Simple skip connection for better gradient flow."""
+    
+    def __init__(self, dim: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+    
+    def forward(self, x):
+        return x + self.norm(x)
+
+
+class NeuralODE(nn.Module):
+    """
+    Wraps a dynamics function with an ODE solver.
+    
+    This lets us:
+    1. Integrate forward in time (simulate the future)
+    2. Backpropagate through the ODE solver (adjoint method)
+    3. Use adaptive step sizing (accuracy + efficiency)
+    """
+
+    def __init__(
+        self,
+        dynamics_fn: DynamicsFunction,
+        solver: str = "dopri5",
+        dt: float = 0.05,
+        rtol: float = 1e-3,
+        atol: float = 1e-4,
+    ):
+        super().__init__()
+        self.dynamics_fn = dynamics_fn
+        self.solver = solver
+        self.dt = dt
+        self.rtol = rtol
+        self.atol = atol
+    
+    def forward(
+        self,
+        initial_state: torch.Tensor,
+        actions: torch.Tensor,
+        time_horizon: float = None,
+    ) -> torch.Tensor:
+        """
+        Roll out the dynamics forward in time.
+        
+        Args:
+            initial_state: [batch, state_dim] starting state
+            actions: [batch, num_steps, action_dim] action sequence
+                     OR [batch, action_dim] for single constant action
+            time_horizon: total simulation time (if None, len(actions) * dt)
+            
+        Returns:
+            trajectory: [batch, num_steps+1, state_dim]
+                       (includes initial state at t=0)
+        """
+        batch_size = initial_state.shape[0]
+        
+        # Handle single action case
+        if actions.dim() == 2:
+            actions = actions.unsqueeze(1)
+        
+        num_steps = actions.shape[1]
+        
+        if time_horizon is None:
+            time_horizon = num_steps * self.dt
+        
+        t_span = torch.linspace(0, time_horizon, num_steps + 1, device=initial_state.device)
+        
+        # Interpolate actions at solver timesteps
+        def dynamics_with_action(t, state):
+            # Find the closest action for current time
+            action_idx = torch.clamp(
+                (t / self.dt).long(), 
+                min=0, 
+                max=num_steps - 1
+            )
+            # Get actions for this batch
+            action = actions[:, action_idx]  # [batch, action_dim]
+            return self.dynamics_fn(state, t, action)
+        
+        # Solve the ODE
+        trajectory = odeint(
+            dynamics_with_action,
+            initial_state,
+            t_span,
+            method=self.solver,
+            rtol=self.rtol,
+            atol=self.atol,
+        )
+        
+        # trajectory shape: [num_timesteps, batch, state_dim]
+        # Transpose to: [batch, num_timesteps, state_dim]
+        trajectory = trajectory.transpose(0, 1)
+        
+        return trajectory
+    
+    def single_step(
+        self,
+        state: torch.Tensor,
+        action: torch.Tensor,
+        dt: float = None,
+    ) -> torch.Tensor:
+        """
+        Single-step prediction: state(t) → state(t+dt)
+        
+        Faster than full rollout for training.
+        """
+        if dt is None:
+            dt = self.dt
+        
+        t_span = torch.tensor([0.0, dt], device=state.device)
+        
+        def dynamics(t, s):
+            return self.dynamics_fn(s, t, action)
+        
+        # Solve for one step
+        result = odeint(
+            dynamics,
+            state,
+            t_span,
+            method="euler",  # fast for single step
+            options={"step_size": dt},
+        )
+        
+        # Return final state: [batch, state_dim]
+        return result[-1]

atlas/core/encoder.py

@@ -0,0 +1,120 @@
+"""
+Encoder: Observation → Latent State
+
+Compresses raw observations (images, vectors, etc.) into
+a compact latent representation that captures the ESSENCE
+of the current state.
+
+Uses variational inference to output a probability distribution
+over possible states, not just a single point. This lets the
+model express UNCERTAINTY about what it's seeing.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Encoder(nn.Module):
+    """
+    Encodes observations into a latent state distribution.
+    
+    Input:  observation tensor [batch, obs_dim]
+    Output: (mean, log_variance) each [batch, latent_dim]
+    
+    The latent state is sampled via the reparameterization trick:
+        z = mean + std * epsilon, where epsilon ~ N(0, 1)
+    
+    This allows gradients to flow through the sampling process.
+    """
+
+    def __init__(self, obs_dim: int, latent_dim: int = 256, hidden_dims: list = None, dropout: float = 0.1):
+        super().__init__()
+        
+        self.obs_dim = obs_dim
+        self.latent_dim = latent_dim
+        
+        if hidden_dims is None:
+            hidden_dims = [512, 512]
+        
+        # Build encoder network
+        layers = []
+        prev_dim = obs_dim
+        
+        for h_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(prev_dim, h_dim),
+                nn.LayerNorm(h_dim),
+                nn.SiLU(),
+                nn.Dropout(dropout),
+            ])
+            prev_dim = h_dim
+        
+        self.feature_net = nn.Sequential(*layers)
+        
+        # Output heads: mean and log_variance of latent distribution
+        self.mean_head = nn.Linear(prev_dim, latent_dim)
+        self.log_var_head = nn.Linear(prev_dim, latent_dim)
+        
+        # Initialize heads with small weights for stable training
+        nn.init.xavier_uniform_(self.mean_head.weight, gain=0.1)
+        nn.init.zeros_(self.mean_head.bias)
+        nn.init.xavier_uniform_(self.log_var_head.weight, gain=0.1)
+        nn.init.zeros_(self.log_var_head.bias)
+    
+    def forward(self, observation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Encode observation into latent distribution parameters.
+        
+        Args:
+            observation: [batch, obs_dim] raw observation
+            
+        Returns:
+            (mean, log_variance): each [batch, latent_dim]
+        """
+        features = self.feature_net(observation)
+        mean = self.mean_head(features)
+        log_var = self.log_var_head(features)
+        
+        # Clamp log_var to prevent numerical instability
+        # Range: [-10, 2] → std range: [~0.00005, ~7.4]
+        log_var = torch.clamp(log_var, min=-10.0, max=2.0)
+        
+        return mean, log_var
+    
+    def sample(self, mean: torch.Tensor, log_var: torch.Tensor) -> torch.Tensor:
+        """
+        Sample from the latent distribution using reparameterization trick.
+        
+        z = μ + σ * ε,  where ε ~ N(0, I)
+        
+        This is differentiable — gradients flow through the sample
+        back to the encoder parameters.
+        """
+        std = torch.exp(0.5 * log_var)  # convert log_var to std
+        epsilon = torch.randn_like(std)  # random noise ~ N(0, 1)
+        z = mean + std * epsilon
+        return z
+    
+    def encode(self, observation: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Full encoding: observation → (mean, log_var, sampled_state)
+        
+        Convenience method that does everything in one call.
+        """
+        mean, log_var = self.forward(observation)
+        z = self.sample(mean, log_var)
+        return mean, log_var, z
+    
+    def kl_divergence(self, mean: torch.Tensor, log_var: torch.Tensor) -> torch.Tensor:
+        """
+        KL divergence from the latent distribution to N(0, I).
+        
+        KL(q(z|x) || p(z)) = -0.5 * Σ(1 + log(σ²) - μ² - σ²)
+        
+        This regularizes the latent space to stay close to a standard
+        normal distribution. Prevents the model from cheating by
+        encoding everything into degenerate distributions.
+        """
+        kl = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp(), dim=-1)
+        return kl.mean()

atlas/core/surprise.py

@@ -0,0 +1,145 @@
+"""
+Surprise Detection System
+
+The KEY insight of the world model: only learn from SURPRISES.
+
+If the model predicted reality correctly → no learning needed
+If the model was WRONG → that's a surprise → UPDATE THE MODEL
+
+This is how the brain works:
+- You expect the floor to be there → you step on it → no surprise
+- You expect the floor to be there → it's not → HUGE surprise → learning!
+
+The surprise signal drives ALL learning in the model.
+This is more efficient than backpropagating through everything —
+we only allocate compute to things we got wrong.
+"""
+
+import torch
+import torch.nn.functional as F
+from collections import deque
+
+
+class SurpriseDetector:
+    """
+    Detects when the world model's predictions don't match reality.
+    
+    Maintains an adaptive threshold:
+    - Gets tighter as the model improves (higher standards)
+    - Loosens when the environment changes (new things to learn)
+    """
+
+    def __init__(self, initial_threshold: float = 0.1, adaptation_rate: float = 0.01, history_size: int = 1000):
+        self.threshold = initial_threshold
+        self.adaptation_rate = adaptation_rate
+        self.history = deque(maxlen=history_size)
+        
+        # Track overall learning progress
+        self.total_surprises = 0
+        self.total_predictions = 0
+        self.surprise_rate_history = deque(maxlen=100)
+    
+    def compute_surprise(self, real: torch.Tensor, predicted: torch.Tensor) -> tuple[torch.Tensor, bool]:
+        """
+        Compute how surprising the real observation is.
+        
+        Args:
+            real: actual observation [batch, obs_dim]
+            predicted: model's prediction [batch, obs_dim]
+            
+        Returns:
+            (surprise_score, is_surprising): scalar loss + boolean
+        """
+        # Element-wise prediction error
+        error = F.mse_loss(real, predicted, reduction="none").mean(dim=-1)  # [batch]
+        
+        # Overall surprise score (mean across batch)
+        surprise_score = error.mean()
+        
+        # Track history
+        self.history.append(surprise_score.item())
+        self.total_predictions += real.shape[0]
+        
+        # Is this surprising enough to trigger learning?
+        is_surprising = surprise_score.item() > self.threshold
+        
+        if is_surprising:
+            self.total_surprises += real.shape[0]
+        
+        # Update adaptive threshold
+        self._adapt_threshold()
+        
+        # Track surprise rate
+        if self.total_predictions > 0:
+            current_rate = self.total_surprises / self.total_predictions
+            self.surprise_rate_history.append(current_rate)
+        
+        return surprise_score, is_surprising
+    
+    def _adapt_threshold(self):
+        """
+        Adapt the surprise threshold based on recent history.
+        
+        Logic:
+        - If surprise rate is high (>50%): we're learning a lot, keep threshold
+        - If surprise rate is low (<10%): model is good, tighten threshold
+        - If surprise rate is medium: maintain current threshold
+        """
+        if len(self.history) < 50:
+            return  # not enough data
+        
+        recent_surprises = list(self.history)[-50:]
+        avg_surprise = sum(recent_surprises) / len(recent_surprises)
+        
+        # Adaptive: threshold follows average surprise but stays slightly below
+        # This means ~50% of predictions will be "surprising"
+        target = avg_surprise * 0.8
+        
+        # Smoothly adjust toward target
+        self.threshold += self.adaptation_rate * (target - self.threshold)
+        
+        # Keep threshold in reasonable bounds
+        self.threshold = max(0.001, min(self.threshold, 10.0))
+    
+    def compute_novelty(self, real: torch.Tensor, predicted: torch.Tensor) -> torch.Tensor:
+        """
+        Compute novelty score — how DIFFERENT is this from what we've seen?
+        
+        High novelty = new situation = important to learn from
+        Low novelty = familiar situation = can skip learning
+        """
+        error = F.mse_loss(real, predicted, reduction="none").mean(dim=-1)
+        
+        # Normalize by recent average surprise
+        if len(self.history) > 0:
+            avg = sum(self.history) / len(self.history)
+            novelty = error / (avg + 1e-8)
+        else:
+            novelty = error
+        
+        return novelty
+    
+    def get_stats(self) -> dict:
+        """Get current surprise statistics."""
+        return {
+            "threshold": self.threshold,
+            "total_surprises": self.total_surprises,
+            "total_predictions": self.total_predictions,
+            "surprise_rate": (
+                self.total_surprises / self.total_predictions
+                if self.total_predictions > 0
+                else 0.0
+            ),
+            "avg_surprise_recent": (
+                sum(self.history) / len(self.history)
+                if len(self.history) > 0
+                else 0.0
+            ),
+        }
+    
+    def reset(self):
+        """Reset all tracking state."""
+        self.history.clear()
+        self.total_surprises = 0
+        self.total_predictions = 0
+        self.surprise_rate_history.clear()