tinysim 0.0.1__py3-none-any.whl

tinysim/__init__.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+
+
+class SimEnvironment(ABC):
+    @abstractmethod
+    def step(self, action) -> dict:
+        pass
+
+    @abstractmethod
+    def reset(self) -> dict:
+        pass

tinysim/_tk_base.py

@@ -0,0 +1,54 @@
+import tkinter as tk
+import threading
+from abc import ABC, abstractmethod
+
+
+class TkBaseFrontend(ABC):
+
+    def __init__(self):
+        self._root = None
+        self._canvas = None
+        self._thread = None
+
+    def render(self):
+        if self._thread is not None:
+            return
+        self._thread = threading.Thread(target=self._window_hook, daemon=True)
+        self._thread.start()
+
+    def _window_hook(self):
+        root = tk.Tk()
+        root.protocol("WM_DELETE_WINDOW", self._on_close)
+        self._create_window(root)
+
+    @abstractmethod
+    def _create_window(self, root):
+        pass
+
+    def bring_to_front(self, root):
+        root.lift()
+        root.attributes("-topmost", True)
+        root.after_idle(root.attributes, "-topmost", False)
+        root.focus_force()
+
+    def _on_close(self):
+        if self._root:
+            try:
+                self._root.destroy()
+            except tk.TclError:
+                pass
+        self._root = None
+        self._canvas = None
+
+    def _pump(self):
+        if not self._root:
+            return
+
+        try:
+            self._root.update_idletasks()
+            self._root.update()
+        except tk.TclError:
+            return
+
+        if self._root:
+            self._root.after(20, self._pump)

tinysim/flappy/__init__.py

@@ -0,0 +1,103 @@
+import numpy as np
+from .. import SimEnvironment
+
+WIDTH, HEIGHT = 800, 600
+GRAVITY = 900.0
+FLAP_STRENGTH = -300.0
+PIPE_SPEED = -200.0
+PIPE_WIDTH = 80
+PIPE_GAP = 200
+PIPE_INTERVAL = 1.6
+
+BIRD_X = 200
+BIRD_SIZE = 35
+
+
+class FlappyEnv(SimEnvironment):
+    def __init__(self, num_envs: int = 1):
+        self.num_envs = num_envs
+        self.reset()
+
+    def reset(self):
+        # birds (vectorized)
+        self.bird_y = np.full(self.num_envs, HEIGHT / 2, dtype=np.float32)
+        self.bird_vel = np.zeros(self.num_envs, dtype=np.float32)
+        self.done = np.zeros(self.num_envs, dtype=bool)
+
+        # shared pipes
+        self.pipes_x = np.empty(0, dtype=np.float32)
+        self.pipes_y = np.empty(0, dtype=np.float32)
+        self.time_since_pipe = PIPE_INTERVAL
+        return self.step(np.zeros(self.num_envs, dtype=np.int32), dt=0.0)
+
+    def _step_physics(self, action, dt):
+        flap_mask = (action == 1) & (~self.done)
+        self.bird_vel[flap_mask] = FLAP_STRENGTH
+        self.bird_vel += GRAVITY * dt
+        self.bird_y += self.bird_vel * dt
+
+    def _spawn_pipe(self):
+        pipe_y = np.random.randint(120, HEIGHT - 120 - PIPE_GAP)
+        self.pipes_x = np.append(self.pipes_x, WIDTH)
+        self.pipes_y = np.append(self.pipes_y, pipe_y)
+
+    def _update_pipes(self, dt):
+        self.time_since_pipe += dt
+        if self.time_since_pipe > PIPE_INTERVAL:
+            self.time_since_pipe = 0.0
+            self._spawn_pipe()
+
+        self.pipes_x += PIPE_SPEED * dt
+
+        # keep pipes on screen
+        keep = self.pipes_x > -PIPE_WIDTH
+        self.pipes_x = self.pipes_x[keep]
+        self.pipes_y = self.pipes_y[keep]
+
+    def _check_collisions(self):  # world bounds
+        hit_bounds = (self.bird_y < 0) | (self.bird_y + BIRD_SIZE > HEIGHT)
+        bx = BIRD_X
+        by = self.bird_y[:, None]
+        px = self.pipes_x[None, :]
+        upper_y = np.zeros_like(self.pipes_y)[None, :]
+        upper_h = self.pipes_y[None, :]
+        lower_y = (self.pipes_y + PIPE_GAP)[None, :]
+        lower_h = (HEIGHT - (self.pipes_y + PIPE_GAP))[None, :]
+        x_overlap = (bx < px + PIPE_WIDTH) & (bx + BIRD_SIZE > px)
+        upper_hit = x_overlap & (by < upper_y + upper_h) & (by + BIRD_SIZE > upper_y)
+        lower_hit = x_overlap & (by < lower_y + lower_h) & (by + BIRD_SIZE > lower_y)
+        hit_pipe = (upper_hit | lower_hit).any(axis=1)
+        self.done |= hit_bounds | hit_pipe
+
+    def step(self, action, dt=0.02):
+        if np.isscalar(action):
+            action = np.full(self.num_envs, action, dtype=np.float32)
+        else:
+            action = np.asarray(action, dtype=np.float32)
+            if action.shape[0] != self.num_envs:
+                raise ValueError(
+                    f"Expected actions of shape ({self.num_envs},), got {action.shape}"
+                )
+
+        # mask done environments to have no action
+        action = action * (~self.done)
+        self._step_physics(action, dt)
+        self._update_pipes(dt)
+        self._check_collisions()
+
+        bird_y = self.bird_y.tolist()
+        bird_vel = self.bird_vel.tolist()
+        done = self.done.tolist()
+
+        if self.num_envs == 1:
+            bird_y = self.bird_y[0]
+            bird_vel = self.bird_vel[0]
+            done = bool(self.done[0])
+
+        return {
+            "bird_y": bird_y,
+            "bird_vel": bird_vel,
+            "pipes_x": self.pipes_x.tolist(),
+            "pipes_y": self.pipes_y.tolist(),
+            "done": done,
+        }

tinysim/flappy/sim.js

@@ -0,0 +1,118 @@
+let simState = {};
+
+export default {
+  initialize({ model }) {
+    model.on("change:sim_state", () => {
+      simState = model.get("sim_state") || {};
+    });
+  },
+
+  async render({ model, el }) {
+    const [width, height] = model.get("_viewport_size") || [800, 600];
+
+    const container = document.createElement("div");
+    container.style.position = "relative";
+    container.style.width = width + "px";
+    container.style.height = height + "px";
+    el.appendChild(container);
+
+    const canvas = document.createElement("canvas");
+    canvas.width = width;
+    canvas.height = height;
+    container.appendChild(canvas);
+    const ctx = canvas.getContext("2d");
+
+    const WORLD_WIDTH = 800;
+    const WORLD_HEIGHT = 600;
+    const BIRD_X = 200;
+    const BIRD_SIZE = 35;
+    const PIPE_WIDTH = 80;
+    const PIPE_GAP = 200;
+    const GROUND_HEIGHT = 80;
+
+    // Scale world to canvas (in case viewport differs from 800×600)
+    const scaleX = width / WORLD_WIDTH;
+    const scaleY = height / WORLD_HEIGHT;
+
+    function draw() {
+      const state = simState || {};
+      const birdY = state.bird_y ?? WORLD_HEIGHT / 2;
+      const pipes_x = state.pipes_x || [];
+      const pipes_y = state.pipes_y || [];
+      const done = state.done || false;
+
+      ctx.clearRect(0, 0, width, height);
+
+      // Background sky
+      ctx.fillStyle = "#70C5CE";
+      ctx.fillRect(0, 0, width, height);
+
+      // Ground
+      const groundH = GROUND_HEIGHT * scaleY;
+      ctx.fillStyle = "#DED895";
+      ctx.fillRect(0, height - groundH, width, groundH);
+
+      // Bird
+      const birdScreenX = BIRD_X * scaleX;
+      const birdScreenY = birdY * scaleY;
+      const birdSizeX = BIRD_SIZE * scaleX;
+      const birdSizeY = BIRD_SIZE * scaleY;
+
+      ctx.fillStyle = "#FFD700";
+      ctx.strokeStyle = "#000000";
+      ctx.lineWidth = 1;
+
+      if (!done) {
+        ctx.beginPath();
+        ctx.ellipse(
+          birdScreenX + birdSizeX / 2,
+          birdScreenY + birdSizeY / 2,
+          birdSizeX / 2,
+          birdSizeY / 2,
+          0,
+          0,
+          Math.PI * 2
+        );
+        ctx.fill();
+        ctx.stroke();
+      }
+
+      // Pipes
+      ctx.fillStyle = "#228B22";
+      ctx.strokeStyle = "#4a8d34";
+      ctx.lineWidth = 2 * ((scaleX + scaleY) / 2);
+
+      for (let i = 0; i < pipes_x.length; i++) {
+        const pipeX = pipes_x[i];
+        const pipeY = pipes_y[i];
+
+        const pw = PIPE_WIDTH * scaleX;
+        const ux = pipeX * scaleX;
+        const uy = 0; // Upper pipe starts at top of screen
+        const uh = pipeY * scaleY; // Upper pipe height extends down to pipeY
+        const lx = pipeX * scaleX;
+        const ly = (pipeY + PIPE_GAP) * scaleY; // Lower pipe starts after gap
+        const lh = (WORLD_HEIGHT - (pipeY + PIPE_GAP)) * scaleY; // Lower pipe extends to bottom
+
+        // Upper pipe
+        ctx.beginPath();
+        ctx.rect(ux, uy, pw, uh);
+        ctx.fill();
+        ctx.stroke();
+
+        // Lower pipe
+        ctx.beginPath();
+        ctx.rect(lx, ly, pw, lh);
+        ctx.fill();
+        ctx.stroke();
+      }
+
+      requestAnimationFrame(draw);
+    }
+
+    draw();
+
+    model.set("_view_ready", true);
+    model.save_changes();
+  }
+};

tinysim/flappy/tk.py

@@ -0,0 +1,94 @@
+import asyncio
+
+try:
+    import tkinter as tk
+    from .. import _tk_base
+except ImportError:
+    raise ImportError("tkinter is required for FlappyTkFrontend")
+
+from . import (
+    FlappyEnv,
+    WIDTH,
+    HEIGHT,
+    PIPE_WIDTH,
+    BIRD_SIZE,
+    BIRD_X,
+    PIPE_GAP,
+)
+
+
+class FlappyTkFrontend(_tk_base.TkBaseFrontend):
+
+    def __init__(self, viewport_size=(800, 600), sim_env=None):
+        super().__init__()
+        if sim_env is None:
+            sim_env = FlappyEnv()
+
+        self.sim_env = sim_env
+        self._viewport_size = viewport_size
+
+    async def step(self, action, dt=0.02):
+        state = self.sim_env.step(action, dt=dt)
+        if self._root:
+            self._root.after(0, lambda: self._draw_state(state))
+
+        await asyncio.sleep(dt)
+        return state
+
+    async def reset(self):
+        state = self.sim_env.reset()
+        if self._root:
+            self._draw_state(state)
+        return state
+
+    def _create_window(self, root):
+        w, h = self._viewport_size
+        root.title("Flappy Bird")
+        canvas = tk.Canvas(root, width=w, height=h, bg="#1E1E1E")
+        canvas.pack(fill="both", expand=True)
+        self._root = root
+        self._canvas = canvas
+        self.bring_to_front(root)
+        self._draw_state(self.sim_env.reset())
+        self._pump()
+        root.mainloop()
+
+    def _draw_state(self, state):
+        if not self._canvas:
+            return
+
+        canvas = self._canvas
+        canvas.delete("all")
+        canvas.create_rectangle(0, 0, WIDTH, HEIGHT, fill="#70C5CE", outline="")
+        canvas.create_rectangle(
+            0, HEIGHT - 80, WIDTH, HEIGHT, fill="#DED895", outline=""
+        )
+
+        # bird
+        by = state["bird_y"]
+
+        if not state.get("done", False):
+            canvas.create_oval(
+                BIRD_X,
+                by,
+                BIRD_X + BIRD_SIZE,
+                by + BIRD_SIZE,
+                fill="#FFD700",
+                outline="#000",
+            )
+
+        # pipes
+        for x, y in zip(state["pipes_x"], state["pipes_y"]):
+            # Upper pipe
+            canvas.create_rectangle(
+                x, 0, x + PIPE_WIDTH, y, fill="#228B22", outline="#4a8d34"
+            )
+            # Lower pipe
+            canvas.create_rectangle(
+                x,
+                y + PIPE_GAP,
+                x + PIPE_WIDTH,
+                HEIGHT,
+                fill="#228B22",
+                outline="#4a8d34",
+            )

tinysim/flappy/widget.py

@@ -0,0 +1,53 @@
+import pathlib
+import anywidget
+import traitlets
+import asyncio
+from IPython.display import display
+from jupyter_ui_poll import ui_events
+
+from . import FlappyEnv
+
+
+class FlappySim(anywidget.AnyWidget):
+    _esm = pathlib.Path(__file__).parent / "sim.js"
+
+    sim_state = traitlets.Dict(default_value={}).tag(sync=True)
+    _viewport_size = traitlets.Tuple(
+        traitlets.Int(), traitlets.Int(), default_value=(800, 600)
+    ).tag(sync=True)
+    _manual_control = traitlets.Bool(default_value=False).tag(sync=True)
+    _view_ready = traitlets.Bool(default_value=False).tag(sync=True)
+
+    def __init__(self, viewport_size=(800, 600), manual_control=False, sim_env=None):
+        super().__init__()
+        self._viewport_size = viewport_size
+        self._manual_control = manual_control
+        if sim_env is None:
+            sim_env = FlappyEnv()
+        if sim_env.num_envs != 1:
+            raise ValueError("FlappySim currently only supports single environment.")
+
+        self.sim_env = sim_env
+        self.sim_state = self.sim_env.reset()
+
+    def render(self):
+        display(self)
+
+        try:
+            with ui_events() as ui_poll:
+                while not self._view_ready:
+                    ui_poll(100)
+        except Exception:
+            pass
+
+    async def step(self, action, dt=0.02):
+        state = self.sim_env.step(action, dt=dt)
+        self.sim_state = state
+        await asyncio.sleep(dt)
+        return state
+
+    async def reset(self):
+        state = self.sim_env.reset()
+        self.sim_state = state
+        await asyncio.sleep(0)
+        return state

tinysim/frogger/__init__.py

@@ -0,0 +1,145 @@
+import numpy as np
+from .. import SimEnvironment
+
+WIDTH, HEIGHT = 800, 600
+CELL = 40
+ROWS, COLS = HEIGHT // CELL, WIDTH // CELL
+
+
+class FroggerEnv(SimEnvironment):
+    def __init__(self, num_envs=1):
+        self.num_envs = num_envs
+        self.total_height = ROWS - 1
+        self.num_cars_per_lane = 4
+        self.car_width = CELL * 2
+        self.traffic_rows = np.array([4, 5, 6, 8, 9])
+        self.speeds = np.array([120, -150, 200, -180, 140], dtype=np.float32)
+        self.reset()
+
+    def reset(self):
+        # Frog positions for all environments (num_envs, 2)
+        self.frog_pos = np.tile(np.array([COLS // 2, ROWS - 1]), (self.num_envs, 1))
+
+        # Cars (num_lanes, num_cars_per_lane)
+        self.car_x = np.zeros(
+            (len(self.traffic_rows), self.num_cars_per_lane), dtype=np.float32
+        )
+        for i in range(self.num_cars_per_lane):
+            self.car_x[:, i] = i * (WIDTH // self.num_cars_per_lane)
+
+        # Score and crossings per environment
+        self.crossings = np.zeros(self.num_envs, dtype=np.float32)
+        self.score = np.zeros(self.num_envs, dtype=np.float32)
+        return self.step(np.zeros(self.num_envs, dtype=np.int32), dt=0.0)
+
+    def _build_car_grid(self):
+        grid = np.zeros((ROWS, COLS), dtype=bool)
+        for lane_idx, row in enumerate(self.traffic_rows):
+            x0 = self.car_x[lane_idx, :]
+            x1 = x0 + self.car_width
+            col_start = np.clip((x0 // CELL).astype(int), 0, COLS - 1)
+            col_end = np.clip((x1 // CELL).astype(int), 0, COLS - 1)
+            for start, end in zip(col_start, col_end):
+                grid[row, start : end + 1] = True
+        return grid
+
+    def step(self, action, dt=0.01):
+        if np.isscalar(action):
+            action = np.full(self.num_envs, action, dtype=np.float32)
+        else:
+            action = np.asarray(action, dtype=np.float32)
+            if action.shape[0] != self.num_envs:
+                raise ValueError(
+                    f"Expected actions of shape ({self.num_envs},), got {action.shape}"
+                )
+        action_map = {
+            0: (0, 0),
+            1: (-1, 0),
+            2: (1, 0),
+            3: (0, -1),
+            4: (0, 1),
+        }
+        dx = np.array([action_map[a][0] for a in action])
+        dy = np.array([action_map[a][1] for a in action])
+
+        # Move frogs
+        self.frog_pos[:, 0] = np.clip(self.frog_pos[:, 0] + dx, 0, COLS - 1)
+        self.frog_pos[:, 1] = np.clip(self.frog_pos[:, 1] + dy, 0, ROWS - 1)
+
+        # Update car positions
+        self.car_x += self.speeds[:, None] * dt
+
+        # Wrap cars
+        pos_mask_pos = self.speeds[:, None] > 0
+        pos_mask_neg = self.speeds[:, None] < 0
+        self.car_x[pos_mask_pos & (self.car_x > WIDTH + 50)] = -self.car_width - 50
+        self.car_x[pos_mask_neg & (self.car_x < -self.car_width - 50)] = WIDTH + 50
+
+        done = np.zeros(self.num_envs, dtype=bool)
+        frog_rects = np.stack(
+            [
+                self.frog_pos[:, 0] * CELL,
+                self.frog_pos[:, 1] * CELL,
+                np.full(self.num_envs, CELL),
+                np.full(self.num_envs, CELL),
+            ],
+            axis=1,
+        )
+
+        self.car_rects = []
+        for lane_idx, row in enumerate(self.traffic_rows):
+            self.car_rects.append(
+                np.stack(
+                    [
+                        self.car_x[lane_idx, :],  # x positions
+                        np.full(self.num_cars_per_lane, row * CELL + 8),  # y positions
+                        np.full(self.num_cars_per_lane, self.car_width),  # widths
+                        np.full(self.num_cars_per_lane, CELL - 16),  # heights
+                    ],
+                    axis=1,
+                )
+            )  # shape: (num_cars_per_lane, 4)
+
+        for env_idx in range(self.num_envs):
+            ax, ay, aw, ah = frog_rects[env_idx]
+            collision = False
+            for car_rects in self.car_rects:
+                bx, by, bw, bh = (
+                    car_rects[:, 0],
+                    car_rects[:, 1],
+                    car_rects[:, 2],
+                    car_rects[:, 3],
+                )
+                overlap = (
+                    (ax < bx + bw) & (ax + aw > bx) & (ay < by + bh) & (ay + ah > by)
+                )
+                if np.any(overlap):
+                    collision = True
+                    break
+            done[env_idx] = collision
+
+        # Handle frogs that reached the top
+        reached_top = self.frog_pos[:, 1] == 0
+        self.crossings[reached_top] += 1
+        self.frog_pos[reached_top] = [COLS // 2, ROWS - 1]
+
+        # Update score
+        current_height = self.total_height - self.frog_pos[:, 1]
+        self.score = self.crossings + current_height / self.total_height
+
+        frog_pos = self.frog_pos.tolist()
+        grid = self._build_car_grid().tolist()
+        scores = self.score.tolist()
+        done = done.tolist()
+
+        if self.num_envs == 1:
+            frog_pos = frog_pos[0]
+            done = done[0]
+            scores = scores[0]
+
+        return {
+            "frog_pos": frog_pos,
+            "grid": grid,
+            "done": done,
+            "score": scores,
+        }

tinysim/frogger/sim.js

@@ -0,0 +1,110 @@
+export default {
+  async render({ model, el }) {
+    const [width, height] = model.get("_viewport_size") || [800, 600];
+
+    const container = document.createElement("div");
+    container.style.position = "relative";
+    container.style.width = width + "px";
+    container.style.height = height + "px";
+    el.appendChild(container);
+
+    const canvas = document.createElement("canvas");
+    canvas.width = width;
+    canvas.height = height;
+    container.appendChild(canvas);
+    const ctx = canvas.getContext("2d");
+
+    // Current sim state
+    let simState = model.get("sim_state") || {};
+    let frogPos = simState.frog_pos || [0, 0];  // grid coords
+    let score = simState.score || 0;
+    let carRects = model.get("car_positions") || []; // pixel absolute coords
+
+    // Watch sim_state and car_positions
+    model.on("change:sim_state", () => {
+      simState = model.get("sim_state") || {};
+      frogPos = simState.frog_pos || frogPos;
+      score = simState.score ?? score;
+    });
+
+    model.on("change:car_positions", () => {
+      carRects = model.get("car_positions") || carRects;
+    });
+
+    const draw = () => {
+      ctx.clearRect(0, 0, width, height);
+
+      // Background
+      ctx.fillStyle = "#101010";
+      ctx.fillRect(0, 0, width, height);
+
+      // The frog is still grid-based
+      const grid = simState.grid || [];
+      const rows = grid.length || 15;
+      const cols = (grid[0] && grid[0].length) || 20;
+
+      const cellW = width / cols;
+      const cellH = height / rows;
+
+      // Safe zones
+      ctx.fillStyle = "#000050";
+      ctx.fillRect(0, 0, width, cellH);
+      ctx.fillStyle = "#004000";
+      ctx.fillRect(0, (rows - 1) * cellH, width, cellH);
+
+      // Road tint
+      for (let r = 1; r < rows - 1; r++) {
+        ctx.fillStyle = "#202020";
+        ctx.fillRect(0, r * cellH, width, cellH);
+      }
+
+      ctx.fillStyle = "#B43232";
+      for (let i = 0; i < carRects.length; i += 4) {
+          ctx.fillRect(carRects[i], carRects[i + 1], carRects[i + 2], carRects[i + 3]);
+      }
+
+      const [frogCol, frogRow] = frogPos;
+      const fx = frogCol * cellW;
+      const fy = frogRow * cellH;
+      const radius = Math.min(cellW, cellH) * 0.4;
+
+      ctx.fillStyle = "#32DC32";
+      ctx.beginPath();
+      ctx.arc(fx + cellW / 2, fy + cellH / 2, radius, 0, Math.PI * 2);
+      ctx.fill();
+
+      // Grid lines
+      ctx.strokeStyle = "#282828";
+      ctx.lineWidth = 1;
+
+      for (let r = 0; r <= rows; r++) {
+        const y = r * cellH;
+        ctx.beginPath();
+        ctx.moveTo(0, y);
+        ctx.lineTo(width, y);
+        ctx.stroke();
+      }
+
+      for (let c = 0; c <= cols; c++) {
+        const x = c * cellW;
+        ctx.beginPath();
+        ctx.moveTo(x, 0);
+        ctx.lineTo(x, height);
+        ctx.stroke();
+      }
+
+      // Score
+      ctx.fillStyle = "#FFFFFF";
+      ctx.font = "16px Arial";
+      ctx.textBaseline = "top";
+      ctx.fillText(`Score: ${score.toFixed(2)}`, 10, 10);
+
+      requestAnimationFrame(draw);
+    };
+
+    draw();
+
+    model.set("_view_ready", true);
+    model.save_changes();
+  },
+};