tinysim 0.0.1__py3-none-any.whl

tinysim/topdown_driving/__init__.py

@@ -0,0 +1,190 @@
+import json
+import math
+import numpy as np
+from pathlib import Path
+from .. import SimEnvironment
+
+
+try:
+    with open(Path(__file__).parent / "track_0.json", "r") as f:
+        track = json.load(f)
+except FileNotFoundError:
+    raise FileNotFoundError("Unable to load track data.")
+
+
+MAX_VEL = 20.0
+ACCELERATION = 8.0
+VEL_FRICT = 2.0
+TURN_SPEED = math.radians(100)
+
+CAR_LENGTH = 1.6  # WORLD units (smaller)
+CAR_WIDTH = 0.8
+CAR_RADIUS = 0.5
+
+LOCAL_WALLS = track["walls"]
+CHECKPOINTS = np.array(track["checkpoints"], dtype=np.float32)
+CHECKPOINT_RADIUS = 3.0
+
+WORLD_WALLS = np.array(
+    [(x, y, w, h, math.radians(rot)) for x, y, w, h, rot in LOCAL_WALLS],
+    dtype=np.float32,
+)
+
+W_X, W_Y, W_W, W_H, W_ROT = WORLD_WALLS.T
+W_HW = W_W * 0.5
+W_HH = W_H * 0.5
+W_COS = np.cos(-W_ROT)
+W_SIN = np.sin(-W_ROT)
+
+# Precompute ray offsets
+RAY_SPREAD = math.radians(75)  # total fan angle
+RAY_LENGTH = 12.0  # world units
+RAY_COUNT = 5
+ray_offsets = np.linspace(
+    -RAY_SPREAD * 0.5, RAY_SPREAD * 0.5, RAY_COUNT, dtype=np.float32
+)
+EPS = 1e-6
+
+
+def cast_rays(x, y, angle):
+    # x,y,angle: (N,)
+    a = angle[:, None] + ray_offsets[None, :]  # (N,R)
+    dx = np.cos(a)
+    dy = np.sin(a)
+    best_t = np.full((x.shape[0], RAY_COUNT), RAY_LENGTH, dtype=np.float32)
+
+    # Expand dims for broadcasting
+    ox = x[:, None, None]
+    oy = y[:, None, None]
+    rdx = dx[:, :, None]
+    rdy = dy[:, :, None]
+
+    wx = W_X[None, None, :]
+    wy = W_Y[None, None, :]
+
+    # transform ray to wall space
+    rox = (ox - wx) * W_COS - (oy - wy) * W_SIN
+    roy = (ox - wx) * W_SIN + (oy - wy) * W_COS
+    rdxl = rdx * W_COS - rdy * W_SIN
+    rdyl = rdx * W_SIN + rdy * W_COS
+
+    tmin = np.full_like(rox, -1e9)
+    tmax = np.full_like(rox, 1e9)
+    mask_x = np.abs(rdxl) > EPS
+    safe_rdxl = np.where(mask_x, rdxl, 1.0)
+    tx1 = (-W_HW - rox) / safe_rdxl
+    tx2 = (W_HW - rox) / safe_rdxl
+    tmin = np.where(mask_x, np.maximum(tmin, np.minimum(tx1, tx2)), tmin)
+    tmax = np.where(mask_x, np.minimum(tmax, np.maximum(tx1, tx2)), tmax)
+    mask_y = np.abs(rdyl) > EPS
+    safe_rdyl = np.where(mask_y, rdyl, 1.0)
+    ty1 = (-W_HH - roy) / safe_rdyl
+    ty2 = (W_HH - roy) / safe_rdyl
+    tmin = np.where(mask_y, np.maximum(tmin, np.minimum(ty1, ty2)), tmin)
+    tmax = np.where(mask_y, np.minimum(tmax, np.maximum(ty1, ty2)), tmax)
+    valid = (tmax >= tmin) & (tmax > 0.0)
+    t = np.where(valid, np.where(tmin > 0, tmin, tmax), np.inf)
+    best_t = np.minimum(best_t, t.min(axis=2))
+    return best_t
+
+
+def collides(cx, cy):
+    dx = cx[:, None] - W_X[None, :]
+    dy = cy[:, None] - W_Y[None, :]
+    lx = dx * W_COS[None, :] - dy * W_SIN[None, :]
+    ly = dx * W_SIN[None, :] + dy * W_COS[None, :]
+    px = np.clip(lx, -W_HW[None, :], W_HW[None, :])
+    py = np.clip(ly, -W_HH[None, :], W_HH[None, :])
+    ddx = lx - px
+    ddy = ly - py
+    hit = (ddx**2 + ddy**2) <= CAR_RADIUS**2
+    return hit.any(axis=1)
+
+
+class TopDownDrivingEnv(SimEnvironment):
+    def __init__(self, num_envs: int = 1):
+        self.num_envs = num_envs
+        self.reset()
+
+    def reset(self):
+        self.x = np.full(self.num_envs, -85.0, dtype=np.float32)
+        self.y = np.full(self.num_envs, -42.0, dtype=np.float32)
+        self.angle = np.zeros(self.num_envs, dtype=np.float32)
+        self.velocity = np.zeros(self.num_envs, dtype=np.float32)
+        self.rays = []
+
+        self.checkpoint_idx = np.zeros(self.num_envs, dtype=np.int32)
+        cx, cy = CHECKPOINTS[0]
+        self.prev_dist = np.hypot(self.x - cx, self.y - cy)
+        return {
+            "x": self.x,
+            "y": self.y,
+            "angle": self.angle,
+            "velocity": self.velocity,
+            "rays": self.rays,
+        }
+
+    def step(self, action, dt=0.02):
+        throttle = action.get("throttle", 0.0)
+        steer = action.get("steer", 0.0)
+
+        if np.isscalar(throttle) and np.isscalar(steer):
+            throttle = np.full(self.num_envs, throttle, dtype=np.float32)
+            steer = np.full(self.num_envs, steer, dtype=np.float32)
+        elif isinstance(throttle, np.ndarray) and isinstance(steer, np.ndarray):
+            throttle = np.asarray(throttle, dtype=np.float32)
+            steer = np.asarray(steer, dtype=np.float32)
+            if throttle.shape[0] != self.num_envs or steer.shape[0] != self.num_envs:
+                raise ValueError(
+                    f"Expected actions of shape ({self.num_envs},), got {throttle.shape} and {steer.shape}"
+                )
+        else:
+            raise ValueError(
+                "Inputs throttle and steer must both be either scalars or numpy arrays."
+            )
+
+        self.velocity += throttle * ACCELERATION * dt
+        self.velocity = np.clip(self.velocity, 0.0, MAX_VEL)
+        self.angle -= steer * TURN_SPEED * dt
+
+        dx = np.cos(self.angle) * self.velocity * dt
+        dy = np.sin(self.angle) * self.velocity * dt
+        nx, ny = self.x + dx, self.y + dy
+
+        hit = collides(nx, ny)
+        self.x = np.where(hit, self.x, nx)
+        self.y = np.where(hit, self.y, ny)
+        self.velocity = np.where(hit, 0.0, self.velocity)
+
+        mask = np.abs(throttle) < 1e-3
+        self.velocity = np.where(
+            mask, np.maximum(0.0, self.velocity - VEL_FRICT * dt), self.velocity
+        )
+        self.rays = cast_rays(self.x, self.y, self.angle)
+
+        # TODO: should reward be scaled by the distance between checkpoints?
+        # return done if the car hits a wall?
+        cp = CHECKPOINTS[self.checkpoint_idx]
+        dist = np.hypot(self.x - cp[:, 0], self.y - cp[:, 1])
+
+        # Reward is the delta to see if the car is getting closer to the checkpoint
+        reward = self.prev_dist - dist
+        reached = dist <= CHECKPOINT_RADIUS
+
+        bonus = 1.0
+        reward += np.where(reached, bonus, 0.0)
+        self.checkpoint_idx = np.where(
+            reached, self.checkpoint_idx + 1, self.checkpoint_idx
+        )
+
+        cp = CHECKPOINTS[self.checkpoint_idx]
+        self.prev_dist = np.hypot(self.x - cp[:, 0], self.y - cp[:, 1])
+
+        return {
+            "x": self.x,
+            "y": self.y,
+            "angle": self.angle,
+            "velocity": self.velocity,
+            "rays": self.rays,
+            "reward": reward,
+        }

tinysim/topdown_driving/sim.js

@@ -0,0 +1,136 @@
+export default {
+  async render({ model, el }) {
+    const [CANVAS_W, CANVAS_H] = model.get("_viewport_size") || [800, 600];
+
+    const container = document.createElement("div");
+    container.style.position = "relative";
+    el.appendChild(container);
+
+    const canvas = document.createElement("canvas");
+    canvas.width = CANVAS_W;
+    canvas.height = CANVAS_H;
+    container.appendChild(canvas);
+
+    const ctx = canvas.getContext("2d");
+
+    const walls = model.get("wall_positions") || [];
+    let sim_state = model.get("sim_state") || {};
+
+    model.on("change:sim_state", () => {
+      sim_state = model.get("sim_state") || {};
+    });
+    const xs = [];
+    const ys = [];
+
+    for (const [x, y, w, h] of walls) {
+      const r = Math.hypot(w, h) * 0.5;
+      xs.push(x - r, x + r);
+      ys.push(y - r, y + r);
+    }
+
+    const min_x = Math.min(...xs);
+    const max_x = Math.max(...xs);
+    const min_y = Math.min(...ys);
+    const max_y = Math.max(...ys);
+
+    const scale = Math.min(
+      (CANVAS_W - 2) / (max_x - min_x),
+      (CANVAS_H - 2) / (max_y - min_y)
+    );
+
+    const offset_x = -min_x * scale;
+    const offset_y = max_y * scale;
+
+    const worldToScreen = (x, y) => {
+      return [
+        x * scale + offset_x,
+        -y * scale + offset_y,
+      ];
+    };
+
+    const drawWalls = () => {
+      ctx.fillStyle = "#cccccc";
+      ctx.strokeStyle = "#000";
+
+      for (const [x, y, w, h, rot] of walls) {
+        const [sx, sy] = worldToScreen(x, y);
+
+        ctx.save();
+        ctx.translate(sx, sy);
+        const rad = (rot * Math.PI) / 180;
+        ctx.rotate(-rad);
+
+        ctx.beginPath();
+        ctx.rect(
+          (-w * scale) / 2,
+          (-h * scale) / 2,
+          w * scale,
+          h * scale
+        );
+        ctx.fill();
+        ctx.stroke();
+
+        ctx.restore();
+      }
+    };
+
+    const drawCars = () => {
+      if (!sim_state.x) return;
+
+      const xs = sim_state.x;
+      const ys = sim_state.y;
+      const angles = sim_state.angle;
+      const CAR_LENGTH = 1.0;
+      const CAR_WIDTH = 0.5;
+      const COLORS = [
+        "red",
+        "orange",
+        "yellow",
+        "green",
+        "blue",
+        "indigo",
+        "violet",
+      ];
+
+      for (let i = 0; i < xs.length; i++) {
+        const [sx, sy] = worldToScreen(xs[i], ys[i]);
+
+        ctx.save();
+        ctx.translate(sx, sy);
+        ctx.rotate(angles[i]);
+
+        ctx.fillStyle = COLORS[i % COLORS.length];
+        ctx.strokeStyle = "black";
+
+        ctx.beginPath();
+        ctx.rect(
+          (-CAR_LENGTH * scale) / 2,
+          (-CAR_WIDTH * scale) / 2,
+          CAR_LENGTH * scale,
+          CAR_WIDTH * scale
+        );
+        ctx.fill();
+        ctx.stroke();
+
+        ctx.restore();
+      }
+    };
+
+    const draw = () => {
+      ctx.clearRect(0, 0, canvas.width, canvas.height);
+
+      ctx.fillStyle = "#ffffff";
+      ctx.fillRect(0, 0, canvas.width, canvas.height);
+
+      drawWalls();
+      drawCars();
+
+      requestAnimationFrame(draw);
+    };
+
+    draw();
+
+    model.set("_view_ready", true);
+    model.save_changes();
+  }
+};

tinysim/topdown_driving/tk.py

@@ -0,0 +1,180 @@
+import asyncio
+import math
+from . import (
+    TopDownDrivingEnv,
+    CAR_LENGTH,
+    CAR_WIDTH,
+    LOCAL_WALLS,
+    CHECKPOINTS,
+    RAY_COUNT,
+    RAY_SPREAD,
+)
+
+try:
+    import tkinter as tk
+    from .. import _tk_base
+except ImportError:
+    raise ImportError("tkinter is required for MountainCarTkFrontend")
+
+
+CHECKPOINT_RADIUS = 0.85
+COLOR_MAP = ["red", "orange", "yellow", "green", "blue", "indigo", "violet"]
+
+xs, ys = [], []
+
+for x, y, w, h, rot in LOCAL_WALLS:
+    r = math.hypot(w, h) * 0.5
+    xs.extend([x - r, x + r])
+    ys.extend([y - r, y + r])
+
+
+min_x, max_x = min(xs), max(xs)
+min_y, max_y = min(ys), max(ys)
+
+CANVAS_W, CANVAS_H = 800, 600
+scale = min((CANVAS_W - 2) / (max_x - min_x), (CANVAS_H - 2) / (max_y - min_y))
+offset_x = -min_x * scale
+offset_y = max_y * scale
+
+
+def rotated_rect(cx, cy, w, h, deg):
+    rad = math.radians(-deg)
+    c, s = math.cos(rad), math.sin(rad)
+    hw, hh = w / 2, h / 2
+    pts = []
+    for x, y in [(-hw, -hh), (hw, -hh), (hw, hh), (-hw, hh)]:
+        rx = x * c - y * s
+        ry = x * s + y * c
+        pts.extend([cx + rx, cy + ry])
+    return pts
+
+
+def world_to_screen(x, y):
+    return x * scale + offset_x, -y * scale + offset_y
+
+
+class TopDownDrivingTkFrontend(_tk_base.TkBaseFrontend):
+
+    def __init__(self, viewport_size=(800, 600), sim_env=None):
+        super().__init__()
+        if sim_env is None:
+            sim_env = TopDownDrivingEnv()
+        self.sim_env = sim_env
+        self._viewport_size = viewport_size
+        self.show_rays = False
+
+        self.keys = set()
+
+    async def step(self, action, dt=0.02):
+        state = self.sim_env.step(action)
+
+        if self._root:
+            self._root.after(0, lambda s=state: self._draw_state(self.sim_env))
+
+        await asyncio.sleep(dt)
+        return state
+
+    async def reset(self):
+        state = self.sim_env.reset()
+        if self._canvas:
+            self._draw_state(self.sim_env)
+        return state
+
+    def _create_window(self, root):
+        w, h = self._viewport_size
+        root.title("Top Down Driving")
+        canvas = tk.Canvas(root, width=w, height=h, bg="white")
+        canvas.pack(expand=True)
+
+        for x, y, w, h, rot in LOCAL_WALLS:
+            cx, cy = world_to_screen(x, y)
+            pts = rotated_rect(cx, cy, w * scale, h * scale, rot)
+            canvas.create_polygon(pts, fill="#cccccc", outline="black")
+
+        r = CHECKPOINT_RADIUS * scale
+        for i, (x, y) in enumerate(CHECKPOINTS):
+            sx, sy = world_to_screen(x, y)
+
+            canvas.create_oval(
+                sx - r,
+                sy - r,
+                sx + r,
+                sy + r,
+                fill="green",
+                outline="",
+                stipple="gray25",
+            )
+            # add checkpoint number
+            canvas.create_text(
+                sx, sy, text=str(i + 1), fill="white", font=("Arial", int(r))
+            )
+
+        root.bind("<KeyPress>", lambda e: self.keys.add(e.keysym))
+        root.bind("<KeyRelease>", lambda e: self.keys.discard(e.keysym))
+
+        self.bring_to_front(root)
+        self._root = root
+        self._canvas = canvas
+        self._draw_state(self.sim_env)
+        self._pump()
+        root.mainloop()
+
+    def _draw_state(self, sim_env):
+        if not self._canvas:
+            return
+
+        c = self._canvas
+        c.delete("car")
+        c.delete("ray")
+
+        xs = sim_env.x
+        ys = sim_env.y
+        angles = sim_env.angle
+        n = len(xs)
+
+        for i in range(n):
+            cx, cy = world_to_screen(xs[i], ys[i])
+
+            pts = rotated_rect(
+                cx,
+                cy,
+                CAR_LENGTH * scale,
+                CAR_WIDTH * scale,
+                math.degrees(angles[i]),
+            )
+            c.create_polygon(
+                pts,
+                fill=COLOR_MAP[i % len(COLOR_MAP)],
+                outline="black",
+                tags="car",
+            )
+
+        if not self.show_rays:
+            return
+
+        ray_offsets = [
+            -RAY_SPREAD * 0.5 + RAY_SPREAD * i / (RAY_COUNT - 1)
+            for i in range(RAY_COUNT)
+        ]
+
+        for i in range(n):
+            ox, oy = xs[i], ys[i]
+            base = angles[i]
+
+            sx1, sy1 = world_to_screen(ox, oy)
+
+            for r_idx, dist in enumerate(sim_env.rays[i]):
+                a = base + ray_offsets[r_idx]
+                x2 = ox + math.cos(a) * dist
+                y2 = oy + math.sin(a) * dist
+                sx2, sy2 = world_to_screen(x2, y2)
+
+                c.create_line(
+                    sx1,
+                    sy1,
+                    sx2,
+                    sy2,
+                    fill="red",
+                    width=1,
+                    tags="ray",
+                )