feat(fragile): update to 0.0.44 (#35)

* feat(fragile): update to 0.0.40

 - depends on a PR to fix `get_best_index` so that the best state and rewards are identified properly.

* feat(fragile): update to 0.0.44

* chore: drop history tree and update r/d scales

- the tree isn't needed for accessing the best attributes
- bring the reward/distance scales back down inline with guillem's recommendations
- suppress the openai retro install print during import of plangym (hacks)

Author: justindujardin

libraries/mathy_python/mathy/agents/fragile.py

@@ -9,6 +9,7 @@
 import numpy as np
 from gym import spaces
 from pydantic import BaseModel
+from wasabi import msg
 
 from mathy import EnvRewards, MathTypeKeysMax, MathyEnvState, is_terminal_transition
 from mathy.envs.gym import MathyGymEnv
@@ -28,15 +29,31 @@ def mathy_dist(x: np.ndarray, y: np.ndarray) -> np.ndarray:
 
 
 def swarm_solve(problem: str, config: SwarmConfig):
-    from plangym import ParallelEnvironment
-    from plangym.env import Environment
-
+    from contextlib import contextmanager
+    import sys, os
+
+    @contextmanager
+    def suppress_stdout():
+        with open(os.devnull, "w") as devnull:
+            old_stdout = sys.stdout
+            sys.stdout = devnull
+            try:
+                yield
+            finally:
+                sys.stdout = old_stdout
+
+    # Suppress the plangym package error "Please install OpenAI retro" that happens
+    # on import. We don't use it, so it's safe to ignore.
+    with suppress_stdout():
+        from plangym import ParallelEnvironment
+        from plangym.core import BaseEnvironment
     from fragile.core.env import DiscreteEnv
     from fragile.core.models import DiscreteModel
     from fragile.core.states import StatesEnv, StatesModel, StatesWalkers
     from fragile.core.swarm import Swarm
     from fragile.core.tree import HistoryTree
     from fragile.core.utils import StateDict
+    from fragile.core.walkers import Walkers
 
     class DiscreteMasked(DiscreteModel):
         def sample(
@@ -69,14 +86,8 @@ def random_choice_prob_index(a, axis=1):
                 **kwargs,
             )
 
-    class MathySwarm(Swarm):
-        def calculate_end_condition(self) -> bool:
-            """Stop when a walker receives a positive terminal reward."""
-            max_reward = self.walkers.env_states.rewards.max()
-            return max_reward > EnvRewards.WIN or self.walkers.calculate_end_condition()
-
     class FragileMathyEnv(DiscreteEnv):
-        """FragileMathyEnv is an interface between the `plangym.Environment` and a
+        """FragileMathyEnv is an interface between the `plangym.BaseEnvironment` and a
         Mathy environment."""
 
         STATE_CLASS = StatesEnv
@@ -120,7 +131,7 @@ def reset(self, batch_size: int = 1, **kwargs) -> StatesEnv:
             )
             return new_states
 
-    class FragileEnvironment(Environment):
+    class FragileEnvironment(BaseEnvironment):
         """Fragile Environment for solving Mathy problems."""
 
         problem: Optional[str]
@@ -216,22 +227,24 @@ def reset(self, return_state: bool = False):
 
     print_every = 1e6
 
-    swarm = MathySwarm(
+    swarm = Swarm(
         model=lambda env: DiscreteMasked(env=env),
         env=lambda: FragileMathyEnv(env=env),
-        tree=HistoryTree,
+        walkers=lambda **kwargs: Walkers(reward_limit=EnvRewards.WIN, **kwargs),
         n_walkers=config.n_walkers,
-        max_iters=config.max_iters,
-        prune_tree=True,
-        reward_scale=5,
-        distance_scale=10,
+        max_epochs=config.max_iters,
+        reward_scale=1,
+        distance_scale=3,
         distance_function=mathy_dist,
         minimize=False,
     )
 
-    _ = swarm.run(print_every=print_every)
-    best_ix = swarm.walkers.states.cum_rewards.argmax()
-    best_id = swarm.walkers.states.id_walkers[best_ix]
-    path = swarm.tree.get_branch(best_id, from_hash=True)
-    last_state = MathyEnvState.from_np(path[0][-1])
-    env._env.mathy.print_history(last_state)
+    with msg.loading(f"Solving {problem} ..."):
+        _ = swarm.run(show_pbar=False)
+
+    if swarm.walkers.best_reward > EnvRewards.WIN:
+        last_state = MathyEnvState.from_np(swarm.walkers.states.best_state)
+        msg.good(f"Solved! {problem} = {last_state.agent.problem}")
+        env._env.mathy.print_history(last_state)
+    else:
+        msg.fail(f"Failed to find a solution! :(")

libraries/mathy_python/requirements.txt

@@ -7,9 +7,7 @@ click
 wasabi
 
 # OpenAI Gym environments
-# 
-# > 0.12.5 has a conflicting version of cloudpickle with TensorFlow 2.0.0
-gym<=0.12.5
+gym
 
 # ML model agents
 tensorflow>=2.1.0

libraries/mathy_python/setup.py

@@ -23,7 +23,7 @@ def setup_package():
         DEVELOPMENT_MODULES = [line.strip() for line in file if "-e" not in line]
 
     extras = {
-        "fragile": ["fragile==0.0.27", "plangym==0.0.2"],
+        "fragile": ["fragile==0.0.44"],
         "dev": DEVELOPMENT_MODULES,
     }
     extras["all"] = [item for group in extras.values() for item in group]