phantomrt 0.1.0__py3-none-any.whl

atlas/environments/physics_2d.py

@@ -0,0 +1,283 @@
+"""
+2D Physics Environment
+
+Simple Newtonian physics simulation with:
+- Balls that move with velocity
+- Gravity
+- Elastic collisions between balls
+- Wall bouncing
+
+This tests whether the world model can learn ACTUAL PHYSICS —
+conservation of momentum, gravity, collision response.
+
+Observation: [ball1_x, ball1_y, ball1_vx, ball1_vy, ball1_r, ball1_color, ...]
+Action: [force_x, force_y] applied to agent ball
+"""
+
+import numpy as np
+from typing import Optional
+from .base import BaseEnvironment
+
+
+class Ball:
+    """Simple physics ball."""
+    
+    def __init__(self, x, y, vx=0, vy=0, radius=0.5, mass=1.0, color=0, fixed=False):
+        self.pos = np.array([x, y], dtype=np.float64)
+        self.vel = np.array([vx, vy], dtype=np.float64)
+        self.radius = radius
+        self.mass = mass
+        self.color = color  # 0=agent, 1=obstacle, 2=collectible
+        self.fixed = fixed  # fixed balls don't move
+
+
+class Physics2DEnvironment(BaseEnvironment):
+    """
+    2D physics world with Newtonian mechanics.
+    
+    Gravity pulls things down, balls bounce off walls
+    and each other, agent can apply forces.
+    """
+
+    def __init__(
+        self,
+        num_balls: int = 4,
+        num_obstacles: int = 2,
+        world_size: float = 10.0,
+        gravity: float = -9.8,
+        dt: float = 0.02,
+        max_steps: int = 500,
+        seed: Optional[int] = None,
+    ):
+        self.num_balls = num_balls
+        self.num_obstacles = num_obstacles
+        self.world_size = world_size
+        self.gravity = gravity
+        self.dt = dt
+        self.max_steps = max_steps
+        
+        self.rng = np.random.RandomState(seed)
+        self.balls = []
+        self.steps = 0
+        
+        # Agent is the first ball (color=0)
+        # Obstacles are static (color=1, fixed)
+        # Other balls are dynamic (color=2)
+        self._obs_dim = num_balls * 6  # x, y, vx, vy, radius, color for each
+        self._action_dim = 2  # force x, force y on agent
+    
+    def get_observation_dim(self) -> int:
+        return self._obs_dim
+    
+    def get_action_dim(self) -> int:
+        return self._action_dim
+    
+    def reset(self) -> np.ndarray:
+        """Reset with random ball positions."""
+        self.steps = 0
+        self.balls = []
+        
+        # Agent ball (red, controllable)
+        agent = Ball(
+            x=self.rng.uniform(1, self.world_size - 1),
+            y=self.rng.uniform(5, self.world_size - 1),
+            vx=0, vy=0,
+            radius=0.5, mass=1.0,
+            color=0, fixed=False,
+        )
+        self.balls.append(agent)
+        
+        # Obstacle balls (green, fixed)
+        for _ in range(self.num_obstacles):
+            obs = Ball(
+                x=self.rng.uniform(1, self.world_size - 1),
+                y=self.rng.uniform(1, 4),
+                radius=0.8, mass=float('inf'),
+                color=1, fixed=True,
+            )
+            self.balls.append(obs)
+        
+        # Dynamic balls (blue)
+        for _ in range(self.num_balls - 1 - self.num_obstacles):
+            ball = Ball(
+                x=self.rng.uniform(1, self.world_size - 1),
+                y=self.rng.uniform(1, self.world_size - 1),
+                vx=self.rng.uniform(-2, 2),
+                vy=self.rng.uniform(-2, 2),
+                radius=0.4, mass=1.0,
+                color=2, fixed=False,
+            )
+            self.balls.append(ball)
+        
+        return self._get_observation()
+    
+    def step(self, action: np.ndarray) -> tuple[np.ndarray, float, bool, dict]:
+        """Apply force and simulate physics."""
+        self.steps += 1
+        
+        # Apply force to agent ball
+        force = np.clip(action, -50, 50).astype(np.float64)
+        agent = self.balls[0]
+        if not agent.fixed:
+            agent.vel += (force / agent.mass) * self.dt
+        
+        # Apply gravity
+        for ball in self.balls:
+            if not ball.fixed:
+                ball.vel[1] += self.gravity * self.dt
+        
+        # Update positions
+        for ball in self.balls:
+            if not ball.fixed:
+                ball.pos += ball.vel * self.dt
+        
+        # Wall collisions
+        for ball in self.balls:
+            if ball.fixed:
+                continue
+            for dim in range(2):
+                if ball.pos[dim] - ball.radius < 0:
+                    ball.pos[dim] = ball.radius
+                    ball.vel[dim] *= -0.8  # bounce with energy loss
+                elif ball.pos[dim] + ball.radius > self.world_size:
+                    ball.pos[dim] = self.world_size - ball.radius
+                    ball.vel[dim] *= -0.8
+        
+        # Ball-ball collisions
+        collisions = 0
+        for i in range(len(self.balls)):
+            for j in range(i + 1, len(self.balls)):
+                if self._collide(self.balls[i], self.balls[j]):
+                    self._resolve_collision(self.balls[i], self.balls[j])
+                    collisions += 1
+        
+        # Compute reward based on agent position relative to other balls
+        reward = 0.0
+        
+        # Small energy penalty (encourages efficient movement)
+        reward -= 0.001 * np.sum(agent.vel ** 2)
+        
+        # Penalty for going out of bounds (shouldn't happen but just in case)
+        if agent.pos[0] < 0 or agent.pos[0] > self.world_size:
+            reward -= 1.0
+        if agent.pos[1] < 0 or agent.pos[1] > self.world_size:
+            reward -= 1.0
+        
+        done = self.steps >= self.max_steps
+        
+        info = {
+            "collisions": collisions,
+            "agent_pos": agent.pos.copy(),
+            "agent_vel": agent.vel.copy(),
+        }
+        
+        return self._get_observation(), reward, done, info
+    
+    def _collide(self, ball1: Ball, ball2: Ball) -> bool:
+        """Check if two balls are overlapping."""
+        dist = np.linalg.norm(ball1.pos - ball2.pos)
+        return dist < (ball1.radius + ball2.radius)
+    
+    def _resolve_collision(self, ball1: Ball, ball2: Ball):
+        """Resolve elastic collision between two balls."""
+        # Separation vector
+        diff = ball1.pos - ball2.pos
+        dist = np.linalg.norm(diff)
+        
+        if dist < 1e-10:
+            return
+        
+        normal = diff / dist
+        
+        # Separate balls
+        overlap = (ball1.radius + ball2.radius) - dist
+        if overlap > 0:
+            if ball1.fixed:
+                ball2.pos -= normal * overlap
+            elif ball2.fixed:
+                ball1.pos += normal * overlap
+            else:
+                ball1.pos += normal * overlap * 0.5
+                ball2.pos -= normal * overlap * 0.5
+        
+        # Compute relative velocity
+        rel_vel = ball1.vel - ball2.vel
+        vel_along_normal = np.dot(rel_vel, normal)
+        
+        # Don't resolve if balls are moving apart
+        if vel_along_normal > 0:
+            return
+        
+        # Elastic collision with restitution
+        restitution = 0.9
+        
+        if ball1.fixed:
+            ball2.vel -= (1 + restitution) * vel_along_normal * normal
+        elif ball2.fixed:
+            ball1.vel += (1 + restitution) * vel_along_normal * normal
+        else:
+            # Standard elastic collision formula
+            m1, m2 = ball1.mass, ball2.mass
+            impulse = -(1 + restitution) * vel_along_normal / (1/m1 + 1/m2)
+            
+            ball1.vel += (impulse / m1) * normal
+            ball2.vel -= (impulse / m2) * normal
+    
+    def _get_observation(self) -> np.ndarray:
+        """Build observation vector."""
+        obs = []
+        
+        # Normalize positions and velocities
+        for ball in self.balls:
+            obs.append(ball.pos[0] / self.world_size)    # x normalized
+            obs.append(ball.pos[1] / self.world_size)    # y normalized
+            obs.append(ball.vel[0] / 10.0)               # vx normalized
+            obs.append(ball.vel[1] / 10.0)               # vy normalized
+            obs.append(ball.radius)                       # radius
+            obs.append(float(ball.color) / 2.0)          # color normalized
+        
+        return np.array(obs, dtype=np.float32)
+    
+    def render(self) -> np.ndarray:
+        """Render physics world as RGB image."""
+        resolution = 100
+        img = np.ones((resolution, resolution, 3), dtype=np.uint8) * 30  # dark background
+        
+        colors = {
+            0: (255, 100, 50),    # agent: orange-red
+            1: (50, 200, 50),     # obstacles: green
+            2: (50, 100, 255),    # dynamic: blue
+        }
+        
+        for ball in self.balls:
+            # Convert world coords to pixel coords
+            px = int(np.clip(ball.pos[0] / self.world_size * (resolution - 1), 0, resolution - 1))
+            py = int(np.clip((1 - ball.pos[1] / self.world_size) * (resolution - 1), 0, resolution - 1))
+            pr = max(1, int(ball.radius / self.world_size * resolution))
+            
+            # Draw ball as filled circle (simplified)
+            y_lo = max(0, py - pr)
+            y_hi = min(resolution, py + pr + 1)
+            x_lo = max(0, px - pr)
+            x_hi = min(resolution, px + pr + 1)
+            
+            color = colors.get(ball.color, (200, 200, 200))
+            img[y_lo:y_hi, x_lo:x_hi] = color
+        
+        return img
+    
+    def get_state(self) -> dict:
+        """Full state for debugging."""
+        return {
+            "balls": [
+                {
+                    "pos": b.pos.copy(),
+                    "vel": b.vel.copy(),
+                    "radius": b.radius,
+                    "color": b.color,
+                    "fixed": b.fixed,
+                }
+                for b in self.balls
+            ],
+            "steps": self.steps,
+        }

atlas/environments/vm_world.py

@@ -0,0 +1,168 @@
+"""
+VMWorldEnv — the real environment: a whole Linux box the agent acts on.
+
+Unlike the toy grid/physics envs, the "action" here is a **shell command string**
+(chosen by the agent from the CommandProposer's candidates). ``step`` runs it in
+the isolated VM and featurizes the *real* outcome into the next observation.
+
+State/observation = features of the last command's real result (exit code, output
+size/entropy, error signatures, timing, …). The world model's job is to predict
+this vector for an action before it runs — prediction error is the surprise signal.
+
+Reward is count-based novelty (intrinsic curiosity): behaviors the agent has rarely
+produced are worth more, so it seeks out unseen corners of the machine. Distinct
+behavior signatures = coverage.
+
+Safety: after a plausibly-destructive or failed command we health-check the VM and,
+if it's bricked, roll back to the base snapshot and continue — recovery, not a
+rewarded action.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import Counter
+from typing import Optional
+
+import numpy as np
+
+from .base import BaseEnvironment
+from ..agents.command_space import embed_command, command_family, ACTION_DIM
+
+STATE_DIM = 24
+
+
+def _entropy(s: str) -> float:
+    if not s:
+        return 0.0
+    counts = Counter(s)
+    n = len(s)
+    return -sum((c / n) * math.log2(c / n) for c in counts.values())
+
+
+class VMWorldEnv(BaseEnvironment):
+    """A real Linux VM presented as an RL-style environment."""
+
+    def __init__(self, vm, log=print, auto_recover: bool = True,
+                 snapshot_tag: str = "base"):
+        self.vm = vm
+        self.log = log
+        self.auto_recover = auto_recover
+        self.snapshot_tag = snapshot_tag
+
+        self.seen = Counter()          # behavior signature -> count (novelty/coverage)
+        self.family_seen = Counter()   # family -> count
+        self._last = np.zeros(STATE_DIM, dtype=np.float32)
+        self.steps = 0
+        self.recoveries = 0
+
+    # ── BaseEnvironment API ──────────────────────────────────────────────────
+    def reset(self) -> np.ndarray:
+        self._last = np.zeros(STATE_DIM, dtype=np.float32)
+        return self._last.copy()
+
+    def get_observation_dim(self) -> int:
+        return STATE_DIM
+
+    def get_action_dim(self) -> int:
+        return ACTION_DIM
+
+    def render(self):
+        return None
+
+    def run_probe(self, command: str, timeout: int = 8) -> np.ndarray:
+        """Run a command purely to measure prediction error (held-out probe);
+        does not mutate learning state."""
+        return self.featurize(command, self.vm.run(command, timeout=timeout))
+
+    # ── the real step ────────────────────────────────────────────────────────
+    def step(self, command: str, timeout: int = 8):
+        """Run ``command`` in the VM; return (obs, reward, done, info) with the
+        real outcome. ``command`` is a string (the agent's chosen action)."""
+        self.steps += 1
+        result = self.vm.run(command, timeout=timeout)
+
+        obs = self.featurize(command, result)
+        sig = self.behavior_signature(command, result)
+
+        # count-based novelty: rarer behavior -> higher intrinsic reward
+        self.seen[sig] += 1
+        self.family_seen[command_family(command)] += 1
+        reward = 1.0 / math.sqrt(self.seen[sig])
+        is_new = self.seen[sig] == 1
+
+        recovered = self._maybe_recover(command, result)
+
+        self._last = obs
+        info = {
+            "command": command,
+            "result": result,
+            "signature": sig,
+            "is_new_behavior": is_new,
+            "coverage": len(self.seen),
+            "recovered": recovered,
+            "family": command_family(command),
+        }
+        return obs, reward, False, info
+
+    # ── featurization: real result -> observation vector ─────────────────────
+    def featurize(self, command: str, result) -> np.ndarray:
+        out, err = result.stdout, result.stderr
+        low = (out + " " + err).lower()
+        v = np.zeros(STATE_DIM, dtype=np.float32)
+        v[0] = max(-1.0, min(1.0, result.exit_code / 128.0))
+        v[1] = 1.0 if result.exit_code == 0 else 0.0
+        v[2] = 1.0 if result.crashed else 0.0
+        v[3] = 1.0 if result.timed_out else 0.0
+        v[4] = min(len(out) / 500.0, 1.0)
+        v[5] = min(len(err) / 500.0, 1.0)
+        v[6] = min(out.count("\n") / 50.0, 1.0)
+        v[7] = _entropy(out[:1000]) / 8.0
+        v[8] = 1.0 if "error" in low else 0.0
+        v[9] = 1.0 if ("not found" in low or "no such" in low) else 0.0
+        v[10] = 1.0 if "permission denied" in low else 0.0
+        v[11] = 1.0 if "usage" in low else 0.0
+        v[12] = 1.0 if any(ch.isdigit() for ch in out[:200]) else 0.0
+        v[13] = 1.0 if out.strip() else 0.0
+        v[14] = 1.0 if err.strip() else 0.0
+        v[15] = min(result.duration / 8.0, 1.0)            # normalized wall time
+        v[16] = embed_command(command)[29]                 # was destructive
+        v[17] = 1.0 if result.exit_code == -1 else 0.0     # launch failure
+        v[18] = 1.0 if result.exit_code == 127 else 0.0    # command not found
+        v[19] = 1.0 if result.exit_code == 126 else 0.0    # not executable
+        v[20] = 1.0 if "\n" in out.strip() else 0.0        # multiline
+        v[21] = 1.0 if len(out) > 2000 else 0.0            # large output
+        v[22] = 1.0 if (result.exit_code == 0 and not out.strip()) else 0.0  # silent success
+        v[23] = 1.0                                        # bias
+        return v
+
+    def behavior_signature(self, command: str, result):
+        """Coarse hash of *what kind of thing happened* — for novelty/coverage."""
+        size_bucket = 0 if not result.stdout else min(len(result.stdout).bit_length(), 12)
+        return (
+            command_family(command),
+            "ok" if result.exit_code == 0 else ("crash" if result.crashed else
+                   ("timeout" if result.timed_out else f"rc{result.exit_code}")),
+            size_bucket,
+            bool(result.stderr.strip()),
+        )
+
+    # ── auto-recovery (the rollback safety net) ──────────────────────────────
+    def _maybe_recover(self, command: str, result) -> bool:
+        if not self.auto_recover:
+            return False
+        # only pay for a health check when something might have broken the box
+        risky = embed_command(command)[29] > 0 or result.exit_code in (-1, 124) \
+            or result.crashed
+        if not risky:
+            return False
+        if self.vm.health_ok():
+            return False
+        self.log(f"[recover] VM unhealthy after: {command!r} → rolling back")
+        try:
+            self.vm.rollback(self.snapshot_tag)
+            self.recoveries += 1
+            return True
+        except Exception as e:
+            self.log(f"[recover] rollback FAILED: {e}")
+            return False

atlas/knowledge/__init__.py

@@ -0,0 +1,3 @@
+from .instruction_vocab import InstructionEncoder, BinaryKnowledgeBase
+
+__all__ = ["InstructionEncoder", "BinaryKnowledgeBase"]