tinysim 0.0.4__py3-none-any.whl → 0.0.5__py3-none-any.whl

tinysim/__init__.py

@@ -1,11 +1,16 @@
 from abc import ABC, abstractmethod
 
+from .tinyspace import Space
+
 
 class SimEnvironment(ABC):
+
+    action_space: Space
+
     @abstractmethod
-    def step(self, action) -> dict:
+    def step(self, *args, **kwargs):
         pass
 
     @abstractmethod
-    def reset(self) -> dict:
+    def reset(self, *args, **kwargs):
         pass

tinysim/_tk_base.py

@@ -1,11 +1,14 @@
-import tkinter as tk
 import threading
+import tkinter as tk
 from abc import ABC, abstractmethod
 
+from . import SimEnvironment
+
 
 class TkBaseFrontend(ABC):
 
-    def __init__(self):
+    def __init__(self, sim_env: SimEnvironment):
+        self.sim_env = sim_env
         self._root = None
         self._canvas = None
         self._thread = None
@@ -16,6 +19,21 @@ class TkBaseFrontend(ABC):
         self._thread = threading.Thread(target=self._window_hook, daemon=True)
         self._thread.start()
 
+    def step(self, *args, **kwargs):
+        state = self.sim_env.step(*args, **kwargs)
+        if self._root:
+            self._root.after(0, lambda: self._draw_state(state))
+        return state
+
+    def reset(self, *args, **kwargs):
+        state = self.sim_env.reset()
+        if self._canvas:
+            self._draw_state(state)
+        return state
+
+    def _draw_state(self, state):
+        raise NotImplementedError
+
     def _window_hook(self):
         root = tk.Tk()
         root.protocol("WM_DELETE_WINDOW", self._on_close)

tinysim/_widget_base.py

@@ -0,0 +1,41 @@
+import anywidget
+import traitlets
+from IPython.display import display
+from jupyter_ui_poll import ui_events
+
+from . import SimEnvironment
+
+
+class BaseWidget(anywidget.AnyWidget):
+
+    sim_env: SimEnvironment
+    _view_ready = traitlets.Bool(default_value=False).tag(sync=True)
+
+    def __init__(self, sim_env: SimEnvironment):
+        super().__init__()
+        self.sim_env = sim_env
+        self.sim_state = self.sim_env.reset()
+        self._update_props(self.sim_state)
+
+    def render(self):
+        display(self)
+
+        try:
+            with ui_events() as ui_poll:
+                while not self._view_ready:
+                    ui_poll(100)
+        except Exception:
+            pass
+
+    def _update_props(self, sim_state):
+        self.sim_state = sim_state
+
+    def step(self, *args, **kwargs):
+        sim_state = self.sim_env.step(*args, **kwargs)
+        self._update_props(sim_state)
+        return sim_state
+
+    def reset(self, *args, **kwargs):
+        sim_state = self.sim_env.reset(*args, **kwargs)
+        self._update_props(sim_state)
+        return sim_state

tinysim/flappy/__init__.py

@@ -1,5 +1,7 @@
 import numpy as np
+
 from .. import SimEnvironment
+from ..tinyspace import Discrete
 
 WIDTH, HEIGHT = 800, 600
 GRAVITY = 900.0
@@ -16,6 +18,7 @@ BIRD_SIZE = 35
 class FlappyEnv(SimEnvironment):
     def __init__(self, num_envs: int = 1):
         self.num_envs = num_envs
+        self.action_space = Discrete(2)  # 0: no flap, 1: flap
         self.reset()
 
     def reset(self):
@@ -54,16 +57,15 @@ class FlappyEnv(SimEnvironment):
         self.pipes_x = self.pipes_x[keep]
         self.pipes_y = self.pipes_y[keep]
 
-    def _check_collisions(self):  # world bounds
+    def _check_collisions(self):
         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)
+        x_overlap = (BIRD_X < px + PIPE_WIDTH) & (BIRD_X + 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)

tinysim/flappy/tk.py

@@ -1,45 +1,22 @@
-import asyncio
-
 try:
     import tkinter as tk
-    from .. import _tk_base
+
+    from .._tk_base import TkBaseFrontend
 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):
+from . import BIRD_SIZE, BIRD_X, HEIGHT, PIPE_GAP, PIPE_WIDTH, WIDTH, FlappyEnv
 
-    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
+class FlappyTkFrontend(TkBaseFrontend):
+    def __init__(self, viewport_size=(800, 600), sim_env=FlappyEnv()):
+        super().__init__(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
+        if sim_env.num_envs != 1:
+            raise ValueError(
+                "FlappyTkFrontend currently only supports single environment."
+            )
 
     def _create_window(self, root):
         w, h = self._viewport_size
@@ -66,7 +43,6 @@ class FlappyTkFrontend(_tk_base.TkBaseFrontend):
 
         # bird
         by = state["bird_y"]
-
         if not state.get("done", False):
             canvas.create_oval(
                 BIRD_X,

tinysim/flappy/widget.py

@@ -1,53 +1,24 @@
 import pathlib
-import anywidget
+
 import traitlets
-import asyncio
-from IPython.display import display
-from jupyter_ui_poll import ui_events
 
+from .._widget_base import BaseWidget
 from . import FlappyEnv
 
 
-class FlappySim(anywidget.AnyWidget):
+class FlappySim(BaseWidget):
     _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)
+    _viewport_size = traitlets.Tuple(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__()
+    def __init__(
+        self, viewport_size=(800, 600), manual_control=False, sim_env=FlappyEnv()
+    ):
+        super().__init__(sim_env)
         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

@@ -1,5 +1,7 @@
 import numpy as np
+
 from .. import SimEnvironment
+from ..tinyspace import Discrete
 
 WIDTH, HEIGHT = 800, 600
 CELL = 40
@@ -14,6 +16,9 @@ class FroggerEnv(SimEnvironment):
         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.done = np.zeros(self.num_envs, dtype=bool)
+
+        self.action_space = Discrete(5)  # 0: stay, 1: left, 2: right, 3: up, 4: down
         self.reset()
 
     def reset(self):
@@ -62,9 +67,17 @@ class FroggerEnv(SimEnvironment):
         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)
+        # Move frogs only move frog if not done
+        self.frog_pos[:, 0] = np.where(
+            ~self.done,
+            np.clip(self.frog_pos[:, 0] + dx, 0, COLS - 1),
+            self.frog_pos[:, 0],
+        )
+        self.frog_pos[:, 1] = np.where(
+            ~self.done,
+            np.clip(self.frog_pos[:, 1] + dy, 0, ROWS - 1),
+            self.frog_pos[:, 1],
+        )
 
         # Update car positions
         self.car_x += self.speeds[:, None] * dt
@@ -127,6 +140,10 @@ class FroggerEnv(SimEnvironment):
         current_height = self.total_height - self.frog_pos[:, 1]
         self.score = self.crossings + current_height / self.total_height
 
+        # Added to persist done state
+        self.done = done | self.done
+        done = self.done
+
         frog_pos = self.frog_pos.tolist()
         grid = self._build_car_grid().tolist()
         scores = self.score.tolist()
@@ -140,6 +157,6 @@ class FroggerEnv(SimEnvironment):
         return {
             "frog_pos": frog_pos,
             "grid": grid,
-            "done": done,
             "score": scores,
+            "done": done,
         }

tinysim/frogger/tk.py

@@ -1,41 +1,24 @@
-import asyncio
-from . import FroggerEnv, WIDTH, HEIGHT, CELL, ROWS, COLS
+from . import CELL, COLS, HEIGHT, ROWS, WIDTH, FroggerEnv
 
 try:
     import tkinter as tk
-    from .. import _tk_base
+
+    from .._tk_base import TkBaseFrontend
 except ImportError:
     raise ImportError("tkinter is required for FroggerTkFrontend")
 
 
-class FroggerTkFrontend(_tk_base.TkBaseFrontend):
+class FroggerTkFrontend(TkBaseFrontend):
+
+    def __init__(self, viewport_size=(800, 600), sim_env=FroggerEnv()):
+        super().__init__(sim_env)
+        self._viewport_size = viewport_size
+        self.keys = set()
 
-    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
-
-        self.keys = set()
-
-    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
@@ -49,11 +32,12 @@ class FroggerTkFrontend(_tk_base.TkBaseFrontend):
         root.bind("<KeyRelease>", lambda e: self.keys.discard(e.keysym))
 
         self.bring_to_front(root)
-        self._draw_state(self.sim_env)
+        self._draw_state()
         self._pump()
         root.mainloop()
 
-    def _draw_state(self, sim_env: FroggerEnv):
+    def _draw_state(self, state=None):
+        sim_env = self.sim_env
         if not self._canvas:
             return
 

tinysim/frogger/widget.py

@@ -1,60 +1,31 @@
 import pathlib
-import anywidget
-import traitlets
-import asyncio
+
 import numpy as np
-from IPython.display import display
-from jupyter_ui_poll import ui_events
+import traitlets
 
+from .._widget_base import BaseWidget
 from . import FroggerEnv
 
 
-class FroggerWidget(anywidget.AnyWidget):
+class FroggerWidget(BaseWidget):
     _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)
+    _viewport_size = traitlets.Tuple(default_value=(800, 600)).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):
+    def __init__(self, viewport_size=(800, 600), sim_env=FroggerEnv()):
+        super().__init__(sim_env)
         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)
+    def _update_props(self, sim_state):
         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

@@ -1,5 +1,7 @@
 import numpy as np
+
 from .. import SimEnvironment
+from ..tinyspace import Discrete
 
 
 class MountainCarEnv(SimEnvironment):
@@ -12,12 +14,12 @@ class MountainCarEnv(SimEnvironment):
         self.force = 0.001
         self.gravity = 0.0025
         self.goal_position = 0.5
+        self.reset()
 
-        # initial state
-        self.position = np.full(num_envs, -0.5, dtype=np.float32)
-        self.velocity = np.zeros(num_envs, dtype=np.float32)
+        # actions: 0 (left), 1 (no push), 2 (right)
+        self.action_space = Discrete(3)
 
-    def step(self, action) -> dict:
+    def step(self, action):
         if np.isscalar(action):
             action = np.full(self.num_envs, action, dtype=np.float32)
         else:
@@ -46,11 +48,14 @@ class MountainCarEnv(SimEnvironment):
 
         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,
-        }
+    def reset(self, **kwargs) -> dict:
+        self.position = np.full(self.num_envs, -0.5, dtype=np.float32)
+        self.velocity = np.zeros(self.num_envs, dtype=np.float32)
+        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}

tinysim/mountain_car/tk.py

@@ -1,39 +1,23 @@
-import asyncio
 import math
+
 import numpy as np
+
 from . import MountainCarEnv
 
 try:
     import tkinter as tk
-    from .. import _tk_base
+
+    from .._tk_base import TkBaseFrontend
 except ImportError:
     raise ImportError("tkinter is required for MountainCarTkFrontend")
 
 
-class MountainCarTkFrontend(_tk_base.TkBaseFrontend):
+class MountainCarTkFrontend(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
+    def __init__(self, viewport_size=(600, 400), sim_env=MountainCarEnv()):
+        super().__init__(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")
@@ -46,7 +30,7 @@ class MountainCarTkFrontend(_tk_base.TkBaseFrontend):
         self._pump()
         root.mainloop()
 
-    def _draw_state(self, state: dict):
+    def _draw_state(self, state=None):
         if not self._canvas:
             return
 
@@ -54,12 +38,10 @@ class MountainCarTkFrontend(_tk_base.TkBaseFrontend):
         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
@@ -105,7 +87,7 @@ class MountainCarTkFrontend(_tk_base.TkBaseFrontend):
         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
+            y_screen = h - y_world * scale
 
             slope = math.cos(3 * x_world)
             angle = -math.atan(slope)

tinysim/mountain_car/widget.py

@@ -1,56 +1,27 @@
 import pathlib
-import anywidget
+
 import traitlets
-import asyncio
-from IPython.display import display
-from jupyter_ui_poll import ui_events
 
+from .._widget_base import BaseWidget
 from . import MountainCarEnv
 
 
-class MountainCarWidget(anywidget.AnyWidget):
+class MountainCarWidget(BaseWidget):
     _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)
+    _viewport_size = traitlets.Tuple(default_value=(600, 400)).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):
+    def __init__(
+        self, manual_control=False, viewport_size=(600, 400), sim_env=MountainCarEnv()
+    ):
+        super().__init__(sim_env)
         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