Merge pull request #12 from FragileTech/alpha

Add skeleton for first release

Author: Guillemdb

pyproject.toml

@@ -59,7 +59,7 @@ dependencies = [
     "quantlib>=1.35",
     "hvplot>=0.11.0",
     "streamz>=0.6.4",
-    "plangym[atari]>=0.1.29",
+    "plangym[dm_control,atari]>=0.1.29",
     "ray>=2.37.0",
     "jupyter",
     "notebook",

src/fragile/actions.py

@@ -0,0 +1,77 @@
+import numpy as np
+import torch
+
+from fragile.core import BaseDtSampler, BaseFractalTree, BasePolicy, FractalTree
+from fragile.fractalai import clone_tensor
+
+
+class UniformDtSampler(BaseDtSampler):
+    def __init__(self, min_dt: int = 1, max_dt: int = 5, fractal: BaseFractalTree | None = None):
+        super().__init__(fractal=fractal)
+        self.max_dt = max_dt
+        self.min_dt = min_dt
+
+    def get_dt(self, n_walkers: int | None = None, fractal: BaseFractalTree | None = None):
+        if n_walkers is None:
+            n_walkers = fractal.n_walkers
+        return np.random.randint(self.min_dt, self.max_dt, size=n_walkers)  # noqa: NPY002
+
+
+class RandomGaussianPolicy(BasePolicy):
+    def __init__(
+        self,
+        std: float = 1.0,
+        min: float | None = None,
+        max: float | None = None,
+        fractal: BaseFractalTree | None = None,
+    ):
+        super().__init__(fractal=fractal)
+        self.std = std
+        self.min_ = min
+        self.max_ = max
+
+    def act(self, n_walkers: int | None = None, fractal: FractalTree | None = None):
+        fractal = fractal if fractal is not None else self.fractal
+        if n_walkers is None:
+            n_walkers = fractal.n_walkers
+        return (torch.randn((n_walkers, *fractal.action_shape)) * self.std).clamp(
+            self.min_, self.max_
+        )
+
+
+class GaussianForce(RandomGaussianPolicy):
+    def __init__(
+        self,
+        std: float = 1.0,
+        min: float | None = None,
+        max: float | None = None,
+        fractal: FractalTree | None = None,
+    ):
+        super().__init__(std=std, fractal=fractal, min=min, max=max)
+        action_shape = fractal.action_shape if fractal is not None else (1,)
+        device = fractal.device if fractal is not None else "cpu"
+        n_walkers = fractal.n_walkers if fractal is not None else 1
+        self.velocity = torch.zeros((n_walkers, *action_shape), device=device)
+
+    def set_fractal(self, fractal: "FractalTree"):
+        super().set_fractal(fractal)
+        self.velocity = torch.zeros(
+            (fractal.n_walkers, *fractal.action_shape), device=fractal.device
+        )
+
+    def act(self, n_walkers: int | None = None, fractal: FractalTree | None = None):
+        action = super().act(n_walkers=n_walkers, fractal=fractal)
+        wc = (
+            self.fractal.will_clone
+            if self.fractal.will_clone.sum() > 0
+            else torch.ones_like(self.fractal.will_clone)
+        )
+        self.velocity[wc] += action
+        return self.velocity[wc].clamp(self.min_, self.max_)
+
+    def clone(self, will_clone: torch.Tensor, clone_ix: torch.Tensor):
+        self.velocity = clone_tensor(self.velocity, clone_ix, will_clone)
+
+    def add_walkers(self, new_walkers):
+        new_vel = torch.zeros((new_walkers, *self.velocity.shape[1:]), device=self.velocity.device)
+        self.velocity = torch.cat((self.velocity, new_vel), dim=0).contiguous()

src/fragile/app/__init__.py

@@ -0,0 +1,64 @@
+import flogging
+import holoviews as hv
+import panel as pn
+
+from fragile.actions import RandomGaussianPolicy, UniformDtSampler
+from fragile.benchmarks import Rastrigin
+from fragile.core import FaiRunner
+from fragile.functions import FunctionTree
+from fragile.shaolin.stream_plots import RGB
+from fragile.shaolin.streaming_fai import InteractiveFai
+
+
+hv.extension("bokeh")
+pn.extension("tabulator", theme="dark")
+
+
+class PlanGymDisplay:
+    def __init__(
+        self,
+    ):
+        self.best_img = RGB()
+        self._curr_best = -1
+
+    def reset(self, fai):  # noqa: ARG002
+        return
+
+    def send(self, fai):
+        best_ix = fai.cum_reward.argmax().cpu().item()
+        best_img = fai.img[best_ix]
+        if best_ix != self._curr_best:
+            self.best_img.send(best_img)
+            self._curr_best = best_ix
+
+    def __panel__(self):
+        return pn.Column(
+            pn.Row(
+                self.best_img.plot,
+                # self.room_grey.plot * self.tree_best_room,
+            ),
+        )
+
+
+def main():
+    flogging.setup()
+    env = Rastrigin(2)
+
+    n_walkers = 10000
+    fai = FunctionTree(
+        max_walkers=n_walkers,
+        env=env,
+        dt_sampler=UniformDtSampler(min_dt=1, max_dt=3),
+        policy=RandomGaussianPolicy(std=0.05, min=-1.0, max=1.0),
+        device="cpu",
+        min_leafs=50,
+        start_walkers=50,
+        minimize=True,
+    )
+    plot = InteractiveFai(fai)
+    runner = FaiRunner(fai, 1000000, plot=plot)
+    return pn.panel(pn.Column(runner, plot)).servable()
+
+
+# if __name__ == "__main__":
+main()

src/fragile/app/debug.py

@@ -0,0 +1,56 @@
+import holoviews as hv
+import panel as pn
+import plangym
+
+from fragile.core import FaiRunner, FractalTree
+from fragile.shaolin.stream_plots import RGB
+
+
+hv.extension("bokeh")
+pn.extension("tabulator", theme="dark")
+
+
+class PlanGymDisplay:
+    def __init__(
+        self,
+    ):
+        self.best_img = RGB()
+        self._curr_best = -1
+
+    def reset(self, fai):  # noqa: ARG002
+        return
+
+    def send(self, fai):
+        best_ix = fai.cum_reward.argmax().cpu().item()
+        best_img = fai.img[best_ix]
+        if best_ix != self._curr_best:
+            self.best_img.send(best_img)
+            self._curr_best = best_ix
+
+    def __panel__(self):
+        return pn.Column(
+            pn.Row(
+                self.best_img.plot,
+                # self.room_grey.plot * self.tree_best_room,
+            ),
+        )
+
+
+def main():
+    env = plangym.make(
+        domain_name="walker",
+        task_name="stand",
+        obs_type="coords",
+        return_image=True,
+        frameskip=1,
+        # n_workers=10,
+        # ray=True,
+    )
+
+    n_walkers = 10000
+    plot = PlanGymDisplay()
+    fai = FractalTree(
+        max_walkers=n_walkers, env=env, device="cpu", min_leafs=250, start_walkers=250
+    )
+    runner = FaiRunner(fai, 1000000, plot=plot)
+    pn.panel(pn.Column(runner, plot)).servable()

src/fragile/app/functions.py

@@ -1,16 +1,13 @@
 from functools import partial
-import threading
-import time
 
 import holoviews as hv
 from holoviews.streams import Pipe
 import numpy as np
-import pandas as pd
 import panel as pn
-import param
 import plangym
 from plangym.utils import process_frame
 
+from fragile.core import FaiRunner
 from fragile.shaolin.stream_plots import Image, RGB
 from fragile.videogames import aggregate_visits, MontezumaTree
 
@@ -19,121 +16,6 @@
 pn.extension("tabulator", theme="dark")
 
 
-class FaiRunner(param.Parameterized):
-    is_running = param.Boolean(default=False)
-
-    def __init__(self, fai, n_steps, plot=None, report_interval=100):
-        super().__init__()
-        self.reset_btn = pn.widgets.Button(icon="restore", button_type="primary")
-        self.play_btn = pn.widgets.Button(icon="player-play", button_type="primary")
-        self.pause_btn = pn.widgets.Button(icon="player-pause", button_type="primary")
-        self.step_btn = pn.widgets.Button(name="Step", button_type="primary")
-        self.progress = pn.indicators.Progress(
-            name="Progress", value=0, width=600, max=n_steps, bar_color="primary"
-        )
-        self.sleep_val = pn.widgets.FloatInput(value=0.0, width=60)
-        self.report_interval = pn.widgets.IntInput(value=report_interval)
-        self.table = pn.widgets.Tabulator()
-        self.fai = fai
-        self.n_steps = n_steps
-        self.curr_step = 0
-        self.plot = plot
-        self.thread = None
-        self.erase_coef_val = pn.widgets.FloatInput(value=0.05, width=60, name="erase")
-
-    @param.depends("erase_coef_val.value")
-    def update_erase_coef(self):
-        self.fai.erase_coef = self.erase_coef_val.value
-
-    @param.depends("reset_btn.value")
-    def on_reset_click(self):
-        self.fai.reset()
-        self.curr_step = 0
-        self.progress.value = 1
-        self.curr_step = 0
-        self.play_btn.disabled = False
-        self.pause_btn.disabled = True
-        self.step_btn.disabled = False
-        self.is_running = False
-        self.progress.bar_color = "primary"
-        summary = pd.DataFrame(self.fai.summary(), index=[0])
-        self.table.value = summary
-        if self.plot is not None:
-            self.plot.reset(self.fai)
-            self.plot.send(self.fai)
-
-    @param.depends("play_btn.value")
-    def on_play_click(self):
-        self.play_btn.disabled = True
-        self.pause_btn.disabled = False
-        self.is_running = True
-        if self.thread is None or not self.thread.is_alive():
-            self.thread = threading.Thread(target=self.run)
-            self.thread.start()
-
-    @param.depends("pause_btn.clicks")
-    def on_pause_click(self):
-        self.is_running = False
-        self.play_btn.disabled = False
-        self.pause_btn.disabled = True
-        if self.thread is not None:
-            self.thread.join()
-
-    @param.depends("step_btn.value")
-    def on_step_click(self):
-        self.take_single_step()
-
-    def take_single_step(self):
-        self.fai.step_tree()
-        self.curr_step += 1
-        self.progress.value = self.curr_step
-        if self.curr_step >= self.n_steps:
-            self.is_running = False
-            self.progress.bar_color = "success"
-            self.step_btn.disabled = True
-            self.play_btn.disabled = True
-            self.pause_btn.disabled = True
-
-        if self.fai.oobs.sum().cpu().item() == self.fai.n_walkers - 1:
-            self.is_running = False
-            self.progress.bar_color = "danger"
-
-        if self.fai.iteration % self.report_interval.value == 0:
-            summary = pd.DataFrame(self.fai.summary(), index=[0])
-            self.table.value = summary
-            if self.plot is not None:
-                self.plot.send(self.fai)
-
-    def run(self):
-        while self.is_running:
-            self.take_single_step()
-            time.sleep(self.sleep_val.value)
-
-    def __panel__(self):
-        # pn.state.add_periodic_callback(self.run, period=20)
-
-        return pn.Column(
-            self.table,
-            self.progress,
-            pn.Row(
-                self.play_btn,
-                self.pause_btn,
-                self.reset_btn,
-                self.step_btn,
-                pn.pane.Markdown("**Sleep**"),
-                self.sleep_val,
-                self.report_interval,
-                self.erase_coef_val,
-            ),
-            self.on_play_click,
-            self.on_pause_click,
-            self.on_reset_click,
-            self.on_step_click,
-            self.update_erase_coef,
-            # self.run,
-        )
-
-
 PYRAMID = [
     [-1, -1, -1, 0, 1, 2, -1, -1, -1],
     [-1, -1, 3, 4, 5, 6, 7, -1, -1],
@@ -324,7 +206,9 @@ def main():
         frameskip=3,
         check_death=True,
         episodic_life=False,
-    )  # , n_workers=10, ray=True)
+        n_workers=10,
+        ray=True,
+    )
 
     n_walkers = 10000
     plot = MontezumaDisplay()
@@ -333,6 +217,3 @@ def main():
     )
     runner = FaiRunner(fai, 1000000, plot=plot)
     pn.panel(pn.Column(runner, plot)).servable()
-
-
-main()