tinysim 0.0.1__py3-none-any.whl

tinysim/frogger/tk.py

@@ -0,0 +1,95 @@
+import asyncio
+from . import FroggerEnv, WIDTH, HEIGHT, CELL, ROWS, COLS
+
+try:
+    import tkinter as tk
+    from .. import _tk_base
+except ImportError:
+    raise ImportError("tkinter is required for FroggerTkFrontend")
+
+
+class FroggerTkFrontend(_tk_base.TkBaseFrontend):
+
+    def __init__(self, viewport_size=(800, 600), sim_env=None):
+        super().__init__()
+        if sim_env is None:
+            sim_env = FroggerEnv()
+        if sim_env.num_envs != 1:
+            raise ValueError(
+                "FroggerTkFrontend currently only supports single environment."
+            )
+        self.sim_env = sim_env
+        self._viewport_size = viewport_size
+
+    async def step(self, action, dt=0.01):
+        state = self.sim_env.step(action, dt=dt)
+        if self._root:
+            self._root.after(0, lambda: 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("Frogger")
+        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)
+        self._pump()
+        root.mainloop()
+
+    def _draw_state(self, sim_env: FroggerEnv):
+        if not self._canvas:
+            return
+
+        canvas = self._canvas
+        canvas.delete("all")
+
+        # safe zones
+        canvas.create_rectangle(0, 0, WIDTH, CELL, fill="#000050", outline="")
+        canvas.create_rectangle(
+            0, (ROWS - 1) * CELL, WIDTH, HEIGHT, fill="#004000", outline=""
+        )
+
+        # cars
+        for lane_rects in sim_env.car_rects:  # lane_rects shape: (num_cars_per_lane, 4)
+            for x, y, w, h in lane_rects:
+                canvas.create_rectangle(x, y, x + w, y + h, fill="#B43232", outline="")
+
+        # frog
+        fx, fy, fw, fh = (
+            sim_env.frog_pos[0, 0] * CELL,
+            sim_env.frog_pos[0, 1] * CELL,
+            CELL,
+            CELL,
+        )
+        canvas.create_oval(
+            fx + 5, fy + 5, fx + fw - 5, fy + fh - 5, fill="#32DC32", outline=""
+        )
+
+        # grid
+        for r in range(ROWS):
+            y = r * CELL
+            canvas.create_line(0, y, WIDTH, y, fill="#282828")
+        for c in range(COLS):
+            x = c * CELL
+            canvas.create_line(x, 0, x, HEIGHT, fill="#282828")
+
+        # score
+        canvas.create_text(
+            10,
+            10,
+            anchor="nw",
+            text=f"Score: {sim_env.score[0]:.2f}",
+            fill="white",
+            font=("Arial", 16),
+        )

tinysim/frogger/widget.py

@@ -0,0 +1,60 @@
+import pathlib
+import anywidget
+import traitlets
+import asyncio
+import numpy as np
+from IPython.display import display
+from jupyter_ui_poll import ui_events
+
+from . import FroggerEnv
+
+
+class FroggerWidget(anywidget.AnyWidget):
+    _esm = pathlib.Path(__file__).parent / "sim.js"
+
+    sim_state = traitlets.Dict(default_value={}).tag(sync=True)
+    car_positions = traitlets.List(default_value=[]).tag(sync=True)
+    _viewport_size = traitlets.Tuple(
+        traitlets.Int(), traitlets.Int(), default_value=(800, 600)
+    ).tag(sync=True)
+    _view_ready = traitlets.Bool(default_value=False).tag(sync=True)
+
+    def get_car_positions(self):
+        return np.vstack(self.sim_env.car_rects).flatten().tolist()
+
+    def __init__(self, viewport_size=(800, 600), sim_env=None):
+        self._viewport_size = viewport_size
+        super().__init__()
+        if sim_env is None:
+            sim_env = FroggerEnv()
+
+        if sim_env.num_envs != 1:
+            raise ValueError(
+                "FroggerWidget currently only supports single environment."
+            )
+
+        self.sim_env = sim_env
+        self.sim_state = self.sim_env.reset()
+        self.car_positions = self.get_car_positions()
+
+    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: int, dt: float = 0.01) -> dict:
+        sim_state = self.sim_env.step(action)
+        self.sim_state = sim_state
+        self.car_positions = self.get_car_positions()
+        await asyncio.sleep(dt)
+        return sim_state
+
+    async def reset(self) -> dict:
+        sim_state = self.sim_env.reset()
+        self.sim_state = sim_state
+        return sim_state

tinysim/mountain_car/__init__.py

@@ -0,0 +1,56 @@
+import numpy as np
+from .. import SimEnvironment
+
+
+class MountainCarEnv(SimEnvironment):
+    def __init__(self, num_envs: int = 1):
+        self.num_envs = num_envs
+        # environment constants
+        self.min_position = -1.2
+        self.max_position = 0.6
+        self.max_speed = 0.07
+        self.force = 0.001
+        self.gravity = 0.0025
+        self.goal_position = 0.5
+
+        # initial state
+        self.position = np.full(num_envs, -0.5, dtype=np.float32)
+        self.velocity = np.zeros(num_envs, dtype=np.float32)
+
+    def step(self, action) -> dict:
+        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}"
+                )
+
+        force_term = (action - 1.0) * self.force
+        self.velocity += force_term + np.cos(3 * self.position) * (-self.gravity)
+        self.velocity = np.clip(self.velocity, -self.max_speed, self.max_speed)
+
+        self.position += self.velocity
+        self.position = np.clip(self.position, self.min_position, self.max_position)
+
+        # handle collision with left wall
+        wall_mask = (self.position == self.min_position) & (self.velocity < 0)
+        self.velocity[wall_mask] = 0.0
+        done = self.position >= self.goal_position
+
+        if self.num_envs == 1:
+            self.position = float(self.position[0])
+            self.velocity = float(self.velocity[0])
+            done = bool(done[0])
+
+        return {"position": self.position, "velocity": self.velocity, "done": done}
+
+    def reset(self) -> dict:
+        self.position = -0.5
+        self.velocity = 0.0
+        return {
+            "position": self.position,
+            "velocity": self.velocity,
+            "done": self.position >= self.goal_position,
+        }

tinysim/mountain_car/sim.js

@@ -0,0 +1,158 @@
+export default {
+  async render({ model, el }) {
+    const [width, height] = model.get("_viewport_size") || [600, 400];
+    const container = document.createElement("div");
+    container.style.position = "relative";
+    el.appendChild(container);
+
+    const canvas = document.createElement("canvas");
+    canvas.width = width;
+    canvas.height = height;
+    container.appendChild(canvas);
+    const ctx = canvas.getContext("2d");
+
+    const sim_state = model.get("sim_state") || {};
+    const minPosition = -1.2;
+    const maxPosition = 0.6;
+    const maxSpeed = 0.07;
+    const force = 0.001;
+    const gravity = 0.0025;
+    let position = sim_state.position || -0.5;
+    let velocity = sim_state.velocity || 0.0;
+
+    if (model.get("_manual_control")) {
+      const controls = document.createElement("div");
+      controls.className = "controls";
+      controls.innerHTML = `
+        <button id="left">←</button>
+        <button id="right">→</button>
+      `;
+
+      const controlsObj = { action: 1 }; // 0=left, 1=idle, 2=right
+      container.appendChild(controls);
+
+      controls.addEventListener("mousedown", (e) => {
+        const id = e.target.id;
+        if (id === "left") {
+          controlsObj.action = 0;
+        } else if (id === "right") {
+          controlsObj.action = 2;
+        }
+      });
+
+      controls.addEventListener("mouseup", (e) => {
+        const id = e.target.id;
+        if (id === "left" || id === "right") {
+          controlsObj.action = 1;
+        }
+      });
+
+      let lastTime = new Date();
+      const step = () => {
+        let ms = new Date();
+        if (ms - lastTime < 20) { 
+          requestAnimationFrame(step);
+          return;
+        }
+        lastTime = ms;
+
+        const action = controlsObj.action; // 0=left, 1=idle, 2=right
+        const forceTerm = (action - 1) * force;
+        velocity += forceTerm + Math.cos(3 * position) * -gravity;
+        velocity = Math.max(Math.min(velocity, maxSpeed), -maxSpeed);
+        position += velocity;
+        position = Math.max(Math.min(position, maxPosition), minPosition);
+        if (position === minPosition && velocity < 0) velocity = 0;
+        requestAnimationFrame(step);
+      };
+
+      step();
+    }
+
+    const goalPosition = 0.5;
+    const worldWidth = maxPosition - minPosition;
+    const scale = width / worldWidth;
+    const heightFn = (x) => Math.sin(3 * x) * 0.45 + 0.55;
+
+    const terrain = [];
+    for (let i = 0; i < 200; i++) {
+      const t = i / 199;
+      const xWorld = minPosition + t * worldWidth;
+      const yWorld = heightFn(xWorld);
+      terrain.push({
+        x: (xWorld - minPosition) * scale,
+        y: height - yWorld * scale
+      });
+    }
+
+    const carWidth = 40;
+    const carHeight = 20;
+    const clearance = 10;
+
+    model.on("change:sim_state", () => {
+      const newState = model.get("sim_state") || {};
+      position = newState.position;
+      velocity = newState.velocity;
+    });
+
+    const draw = () => {
+      ctx.clearRect(0, 0, width, height);
+      ctx.fillStyle = "#eeeeee";
+      ctx.beginPath();
+      ctx.moveTo(terrain[0].x, terrain[0].y);
+      for (const p of terrain) ctx.lineTo(p.x, p.y);
+      ctx.lineTo(width, height);
+      ctx.lineTo(0, height);
+      ctx.closePath();
+      ctx.fill();
+
+      ctx.strokeStyle = "#444";
+      ctx.lineWidth = 2;
+      ctx.stroke();
+      const xScreen = (position - minPosition) * scale;
+      const yScreen = height - heightFn(position) * scale - clearance;
+      const slope = Math.cos(3 * position);
+      const angle = Math.atan(slope);
+
+      ctx.save();
+      ctx.translate(xScreen, yScreen);
+      ctx.rotate(-angle);
+
+      ctx.fillStyle = "#000";
+      ctx.fillRect(-carWidth / 2, -carHeight, carWidth, carHeight);
+
+      ctx.fillStyle = "#777";
+      const wheelR = carHeight * 0.4;
+      ctx.beginPath();
+      ctx.arc(carWidth / 3, 0, wheelR, 0, Math.PI * 2);
+      ctx.fill();
+      ctx.beginPath();
+      ctx.arc(-carWidth / 3, 0, wheelR, 0, Math.PI * 2);
+      ctx.fill();
+      ctx.restore();
+
+      const fx = (goalPosition - minPosition) * scale;
+      const fy = height - heightFn(goalPosition) * scale;
+      ctx.strokeStyle = "#000";
+      ctx.beginPath();
+      ctx.moveTo(fx, fy);
+      ctx.lineTo(fx, fy - 40);
+      ctx.stroke();
+
+      ctx.fillStyle = "#ff0";
+      ctx.beginPath();
+      ctx.moveTo(fx, fy - 40);
+      ctx.lineTo(fx + 25, fy - 35);
+      ctx.lineTo(fx, fy - 30);
+      ctx.closePath();
+      ctx.fill();
+
+      requestAnimationFrame(draw);
+    };
+
+    draw();
+
+    model.set("_view_ready", true);
+    model.save_changes();
+  }
+};

tinysim/mountain_car/styles.css

@@ -0,0 +1,5 @@
+.controls {
+  display: flex;
+  margin-top: 10px;
+  gap: 10px;
+}

tinysim/mountain_car/tk.py

@@ -0,0 +1,141 @@
+import asyncio
+import math
+import numpy as np
+from . import MountainCarEnv
+
+try:
+    import tkinter as tk
+    from .. import _tk_base
+except ImportError:
+    raise ImportError("tkinter is required for MountainCarTkFrontend")
+
+
+class MountainCarTkFrontend(_tk_base.TkBaseFrontend):
+
+    def __init__(self, viewport_size=(600, 400), sim_env=None):
+        super().__init__()
+        if sim_env is None:
+            sim_env = MountainCarEnv()
+        self.sim_env = sim_env
+        self._viewport_size = viewport_size
+
+    async def step(self, action, dt=0.01):
+        state = self.sim_env.step(action)
+
+        if self._root:
+            self._root.after(0, lambda s=state: self._draw_state(s))
+
+        await asyncio.sleep(dt)
+        return state
+
+    async def reset(self):
+        state = self.sim_env.reset()
+        if self._canvas:
+            self._draw_state(state)
+        return state
+
+    def _create_window(self, root):
+        w, h = self._viewport_size
+        root.title("Mountain Car")
+        canvas = tk.Canvas(root, width=w, height=h, bg="#eeeeee")
+        canvas.pack(fill="both", expand=True)
+        self.bring_to_front(root)
+        self._root = root
+        self._canvas = canvas
+        self._draw_state(self.sim_env.reset())
+        self._pump()
+        root.mainloop()
+
+    def _draw_state(self, state: dict):
+        if not self._canvas:
+            return
+
+        c = self._canvas
+        w = int(c.winfo_width() or self._viewport_size[0])
+        h = int(c.winfo_height() or self._viewport_size[1])
+        c.delete("all")
+
+        min_x = self.sim_env.min_position
+        max_x = self.sim_env.max_position
+        world_width = max_x - min_x
+        scale = w / world_width
+        clearance = 10
+
+        def heightFn(x):
+            return np.sin(3 * x) * 0.45 + 0.55
+
+        # Terrain
+        terrain_pts = []
+        for px in range(w):
+            x_world = min_x + px / w * world_width
+            y_world = heightFn(x_world)
+            y_screen = h - y_world * scale
+            terrain_pts.append((px, y_screen))
+        for i in range(len(terrain_pts) - 1):
+            c.create_line(*terrain_pts[i], *terrain_pts[i + 1], fill="#444444", width=2)
+
+        # Goal
+        gx = self.sim_env.goal_position
+        gy = heightFn(gx)
+        goal_x = (gx - min_x) * scale
+        goal_y = h - gy * scale
+        c.create_line(goal_x, goal_y, goal_x, goal_y - 40, fill="#000", width=2)
+        c.create_polygon(
+            goal_x,
+            goal_y - 40,
+            goal_x + 25,
+            goal_y - 35,
+            goal_x,
+            goal_y - 30,
+            fill="#ffff00",
+            outline="",
+        )
+
+        # Draw each car
+        positions = np.atleast_1d(state["position"])
+        car_width, car_height = 40, 20
+        wheel_r = car_height * 0.4
+
+        def rot(px, py, ang):
+            s, c0 = math.sin(ang), math.cos(ang)
+            return px * c0 - py * s, px * s + py * c0
+
+        colors = ["#ff0000", "#00aa00", "#0000ff", "#ffaa00"]
+
+        for i, x_world in enumerate(positions):
+            y_world = heightFn(x_world)
+            x_screen = (x_world - min_x) * scale
+            y_screen = h - y_world * scale - clearance
+
+            slope = math.cos(3 * x_world)
+            angle = -math.atan(slope)
+
+            # body
+            body_local = [
+                (-car_width / 2, -car_height),
+                (car_width / 2, -car_height),
+                (car_width / 2, 0),
+                (-car_width / 2, 0),
+            ]
+            body_screen = []
+            for px, py in body_local:
+                rx, ry = rot(px, py, angle)
+                body_screen.append((x_screen + rx, y_screen + ry))
+
+            c.create_polygon(
+                body_screen, fill=colors[i % len(colors)], outline="", stipple="gray75"
+            )
+
+            # wheels
+            for wx in (-car_width / 3, car_width / 3):
+                rx, ry = rot(wx, 0, angle)
+                cx = x_screen + rx
+                cy = y_screen + ry
+                c.create_oval(
+                    cx - wheel_r,
+                    cy - wheel_r,
+                    cx + wheel_r,
+                    cy + wheel_r,
+                    fill="#777777",
+                    outline="",
+                )

tinysim/mountain_car/widget.py

@@ -0,0 +1,56 @@
+import pathlib
+import anywidget
+import traitlets
+import asyncio
+from IPython.display import display
+from jupyter_ui_poll import ui_events
+
+from . import MountainCarEnv
+
+
+class MountainCarWidget(anywidget.AnyWidget):
+    _esm = pathlib.Path(__file__).parent / "sim.js"
+    _css = pathlib.Path(__file__).parent / "styles.css"
+
+    sim_state = traitlets.Dict(default_value={}).tag(sync=True)
+    _viewport_size = traitlets.Tuple(
+        traitlets.Int(), traitlets.Int(), default_value=(600, 400)
+    ).tag(sync=True)
+    _view_ready = traitlets.Bool(default_value=False).tag(sync=True)
+    _manual_control = traitlets.Bool(default_value=False).tag(sync=True)
+
+    def __init__(self, manual_control=False, viewport_size=(600, 400), sim_env=None):
+        self._manual_control = manual_control
+        self._viewport_size = viewport_size
+        super().__init__()
+        if sim_env is None:
+            sim_env = MountainCarEnv()
+
+        if sim_env.num_envs != 1:
+            raise ValueError(
+                "MountainCarWidget 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: int, dt: float = 0.01) -> dict:
+        sim_state = self.sim_env.step(action)
+        self.sim_state = sim_state
+        await asyncio.sleep(dt)
+        return sim_state
+
+    async def reset(self) -> dict:
+        sim_state = self.sim_env.reset()
+        self.sim_state = sim_state
+        return sim_state

tinysim/simple_amr/example_maps.py

@@ -0,0 +1,121 @@
+import random
+import math
+import copy
+
+empty_map = {
+    "name": "rect_triangle_map",
+    "robot": {
+        "pos": [500, 300],
+        "angle": 0,
+    },  # speed: 0.01, turn_speed: 0.04, lidar: { numBeams: 30, fov: 6.28, maxRange: 1000}
+    "map": [
+        {
+            "type": "rectangle",
+            "x": 400,
+            "y": 0,
+            "width": 800,
+            "height": 30,
+            "bodyInfo": {"angle": 0, "render": {"fillStyle": "#2f6f8f"}},
+            "label": "Wall",
+        },
+        {
+            "type": "rectangle",
+            "x": 400,
+            "y": 590,
+            "width": 800,
+            "height": 20,
+            "bodyInfo": {"angle": 0, "render": {"fillStyle": "#2f6f8f"}},
+            "label": "Wall",
+        },
+        {
+            "type": "rectangle",
+            "x": 10,
+            "y": 500,
+            "width": 20,
+            "height": 1000,
+            "bodyInfo": {"angle": 0, "render": {"fillStyle": "#2f6f8f"}},
+            "label": "Wall",
+        },
+        {
+            "type": "rectangle",
+            "x": 790,
+            "y": 500,
+            "width": 20,
+            "height": 1000,
+            "bodyInfo": {"angle": 0, "render": {"fillStyle": "#2f6f8f"}},
+            "label": "Wall",
+        },
+    ],
+}
+
+
+def gen_simple_map():
+    map_data = copy.deepcopy(empty_map)
+    width = 800
+    height = 600
+
+    # Generate static obstacles
+    for _ in range(6):
+        w = 60 + random.random() * 120
+        h = 40 + random.random() * 100
+        x = 120 + random.random() * (width - 240)
+        y = 120 + random.random() * (height - 240)
+        angle = random.random() * math.pi
+        rect = {
+            "type": "rectangle",
+            "x": x,
+            "y": y,
+            "width": w,
+            "height": h,
+            "bodyInfo": {
+                "angle": angle,
+                "render": {"fillStyle": "#2f6f8f"},
+                "label": "Obstacle",
+            },
+        }
+        map_data["map"].append(rect)
+
+    # Generate dynamic pushable boxes
+    for _ in range(3):
+        size = 40 + random.random() * 20
+        x = 100 + random.random() * (width - 200)
+        y = 100 + random.random() * (height - 200)
+        box = {
+            "type": "rectangle",
+            "x": x,
+            "y": y,
+            "width": size,
+            "height": size,
+            "bodyInfo": {
+                "isStatic": False,
+                "frictionAir": 0.1,
+                "friction": 0.5,
+                "restitution": 0.2,
+                "density": 0.002,
+                "render": {"fillStyle": "#b2df8a"},
+                "label": "Pushable Box",
+            },
+        }
+        map_data["map"].append(box)
+
+    for i in range(2):
+        size = 40 + random.random() * 20
+        x = 100 + random.random() * (width - 200)
+        y = 100 + random.random() * (height - 200)
+        angle = random.random() * math.pi
+        box = {
+            "type": "rectangle",
+            "x": x,
+            "y": y,
+            "width": size,
+            "height": size,
+            "bodyInfo": {
+                "isStatic": True,
+                "angle": angle,
+                "label": "Goal",
+                "render": {"fillStyle": "#ffffff"},
+            },
+        }
+        map_data["map"].append(box)
+
+    return map_data