feat(fragile): add swarm agent as untrained solver

justindujardin · justindujardin · commit 85152731b581 · 2020-04-01T13:09:32.000-07:00
- uses the crazy cool FractalAI fragile library to search the mathy env state spaces and find solutions without a trained neural network. - check them out: https://github.com/FragileTech/fragile
diff --git a/libraries/mathy_python/mathy/agents/fragile.py b/libraries/mathy_python/mathy/agents/fragile.py
@@ -1,39 +1,34 @@
-from typing import List
+"""Use Fractal Monte Carlo search in order to solve mathy problems without a
+trained neural network."""
+
 import copy
-import sys
 import time
-import traceback
 
 import gym
-import holoviews
 import numpy as np
-import srsly
 from gym import spaces
+from plangym import ParallelEnvironment
+from plangym.env import Environment
 
 from fragile.core.env import DiscreteEnv
-from fragile.core.models import Bounds, DiscreteModel
+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
-from fragile.dataviz import AtariViz, LandscapeViz, Summary, SwarmViz, SwarmViz1D
-from mathy import MathTypeKeysMax, MathyEnvState, is_terminal_transition, EnvRewards
+from mathy import EnvRewards, MathTypeKeysMax, MathyEnvState, is_terminal_transition
 from mathy.envs.gym import MathyGymEnv
-from plangym import ParallelEnvironment
-from plangym.env import Environment
-import os
+from typing import List, Optional
+
+from pydantic import BaseModel
+from .. import about
 
 
-# Print explored mathy states when True
-verbose = False
-use_mp = True
-prune_tree = True
-max_iters = 100
-reward_scale = 5
-distance_scale = 10
-minimize = False
-use_vis = False
+class SwarmConfig(BaseModel):
+    verbose: bool = False
+    use_mp: bool = True
+    n_walkers: int = 512
+    max_iters: int = 100
 
 
 class DiscreteMasked(DiscreteModel):
@@ -80,23 +75,23 @@ class FragileMathyEnv(DiscreteEnv):
 
     def get_params_dict(self) -> StateDict:
         super_params = super(FragileMathyEnv, self).get_params_dict()
-        params = {"game_ends": {"dtype": np.bool_}}
+        params = {"terminals": {"dtype": np.bool_}}
         params.update(super_params)
         return params
 
     def step(self, model_states: StatesModel, env_states: StatesEnv) -> StatesEnv:
         actions = model_states.actions.astype(np.int32)
-        new_states, observs, rewards, game_ends, infos = self._env.step_batch(
+        new_states, observs, rewards, terminals, infos = self._env.step_batch(
             actions=actions, states=env_states.states
         )
-        ends = [not inf.get("valid", False) for inf in infos]
+        oobs = [not inf.get("valid", False) for inf in infos]
         new_state = self.states_from_data(
             states=new_states,
             observs=observs,
             rewards=rewards,
-            ends=ends,
+            oobs=oobs,
             batch_size=len(actions),
-            game_ends=game_ends,
+            terminals=terminals,
         )
         return new_state
 
@@ -105,40 +100,47 @@ def reset(self, batch_size: int = 1, **kwargs) -> StatesEnv:
         states = np.array([copy.deepcopy(state) for _ in range(batch_size)])
         observs = np.array([copy.deepcopy(obs) for _ in range(batch_size)])
         rewards = np.zeros(batch_size, dtype=np.float32)
-        ends = np.zeros(batch_size, dtype=np.bool_)
-        game_ends = np.zeros(batch_size, dtype=np.bool_)
+        oobs = np.zeros(batch_size, dtype=np.bool_)
+        terminals = np.zeros(batch_size, dtype=np.bool_)
         new_states = self.states_from_data(
             states=states,
             observs=observs,
             rewards=rewards,
-            ends=ends,
+            oobs=oobs,
             batch_size=batch_size,
-            game_ends=game_ends,
+            terminals=terminals,
         )
         return new_states
 
 
 class FragileEnvironment(Environment):
     """Fragile Environment for solving Mathy problems."""
 
+    problem: Optional[str]
+
     def __init__(
         self,
         name: str,
         environment: str = "poly",
         difficulty: str = "normal",
+        problem: str = None,
+        max_steps: int = 64,
         wrappers=None,
         **kwargs,
     ):
         super(FragileEnvironment, self).__init__(name=name)
         self._env: MathyGymEnv = gym.make(
             f"mathy-{environment}-{difficulty}-v0",
-            verbose=verbose,
             np_observation=True,
+            error_invalid=False,
+            env_problem=problem,
             **kwargs,
         )
         self.observation_space = spaces.Box(
             low=0, high=MathTypeKeysMax, shape=(256, 256, 1), dtype=np.uint8,
         )
+        self.problem = problem
+        self.max_steps = max_steps
         self.wrappers = wrappers
         self.init_env()
 
@@ -159,7 +161,7 @@ def __getattr__(self, item):
         return getattr(self._env, item)
 
     def get_state(self) -> np.ndarray:
-        assert self._env is not None, "env required to get_state"
+        assert self._env.state is not None, "env required to get_state"
         return self._env.state.to_np()
 
     def set_state(self, state: np.ndarray):
@@ -238,38 +240,69 @@ def solve_problem(problem_env: str, problem_difficulty: str = "easy"):
     # print_every = 5 if problem_difficulty == "hard" else 10
     print_every = 1e6
     n_walkers = 768 if problem_difficulty == "hard" else 256
+    max_iters = 100
 
     swarm = MathySwarm(
         model=lambda env: DiscreteMasked(env=env),
         env=lambda: FragileMathyEnv(env=env),
         tree=HistoryTree,
         n_walkers=n_walkers,
         max_iters=max_iters,
-        prune_tree=prune_tree,
-        reward_scale=reward_scale,
-        distance_scale=distance_scale,
+        prune_tree=True,
+        reward_scale=5,
+        distance_scale=10,
         distance_function=mathy_dist,
-        minimize=minimize,
+        minimize=False,
     )
 
-    if not use_vis:
-        _ = swarm.run(print_every=print_every)
+    _ = 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)
+
+
+def swarm_solve(problem: str, config: SwarmConfig):
+    if config.use_mp:
+        env = ParallelEnvironment(
+            env_class=FragileEnvironment,
+            name="mathy_v0",
+            problem=problem,
+            repeat_problem=True,
+        )
     else:
-        holoviews.extension("bokeh")
-        viz = SwarmViz1D(swarm, stream_interval=print_every)
-        viz.plot()
-        viz.run(print_every=print_every)
+        env = FragileEnvironment(name="mathy_v0", problem=problem, repeat_problem=True)
 
+    print_every = 1e6
+
+    swarm = MathySwarm(
+        model=lambda env: DiscreteMasked(env=env),
+        env=lambda: FragileMathyEnv(env=env),
+        tree=HistoryTree,
+        n_walkers=config.n_walkers,
+        max_iters=config.max_iters,
+        prune_tree=True,
+        reward_scale=5,
+        distance_scale=10,
+        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)
-    env.render(last_action=-1, last_reward=0.0)
-    for s, a in zip(path[0][1:], path[1]):
-        _, _, r, _, info = env.step(state=s, action=a)
-        env.render(last_action=a, last_reward=r)
-        time.sleep(0.01)
+    last_state = MathyEnvState.from_np(path[0][-1])
+    env._env.mathy.print_history(last_state)
+
 
+if __name__ == "__main__":
+    # Solve one problem of each type and difficulty
+    for e_type in ["complex", "binomial", "poly", "poly-blockers"]:
+        for e_difficulty in ["easy", "normal"]:
+            solve_problem(e_type, e_difficulty)
 
-for e_type in ["complex", "binomial", "poly", "poly-blockers"]:
-    for e_difficulty in ["easy", "normal", "hard"]:
-        solve_problem(e_type, e_difficulty)
+    # # Solve a bunch of problems of the same type
+    # for i in range(10):
+    #     solve_problem("poly", "easy")
diff --git a/libraries/mathy_python/mathy/api.py b/libraries/mathy_python/mathy/api.py
@@ -1,38 +1,66 @@
-from typing import Optional
+from dataclasses import dataclass
+from typing import Optional, Union
 
 from .agents.base_config import BaseConfig
 from .agents.episode_memory import EpisodeMemory
+from .agents.fragile import SwarmConfig, swarm_solve
 from .agents.policy_value_model import PolicyValueModel, load_policy_value_model
+from .state import MathyEnvState
+
+
+@dataclass
+class MathyAPIModelState:
+    config: BaseConfig
+    model: PolicyValueModel
+
+
+@dataclass
+class MathyAPISwarmState:
+    config: SwarmConfig
 
 
 class Mathy:
     """The standard interface for working with Mathy models and agents."""
 
-    config: BaseConfig
-    model: PolicyValueModel
+    state: Union[MathyAPIModelState, MathyAPISwarmState]
 
     def __init__(
         self,
         *,
         model_path: str = None,
         model: PolicyValueModel = None,
-        config: BaseConfig = None,
+        config: Union[BaseConfig, SwarmConfig] = None,
         silent: bool = False,
     ):
         if model_path is not None:
-            self.model, self.config = load_policy_value_model(model_path, silent=silent)
+            model, config = load_policy_value_model(model_path, silent=silent)
+            self.state = MathyAPIModelState(model=model, config=config)
         elif model is not None and config is not None:
             if not isinstance(model, PolicyValueModel):
                 raise ValueError("model must derive PolicyValueModel for compatibility")
-            self.model = model
-            self.config = config
+            if not isinstance(config, BaseConfig):
+                raise ValueError("config must be a BaseConfig instance")
+            self.state = MathyAPIModelState(model=model, config=config)
         else:
-            raise ValueError(
-                "Either 'model_path' or ('model' and 'config') must be provided"
-            )
+            if not isinstance(config, SwarmConfig):
+                raise ValueError("config must be a BaseConfig instance")
+            self.state = MathyAPISwarmState(config=config)
 
     def simplify(
         self, *, model: str = "mathy_alpha_sm", problem: str, max_steps: int = 128,
+    ) -> EpisodeMemory:
+        if isinstance(self.state, MathyAPISwarmState):
+            return self.simplify_swarm(problem=problem, max_steps=max_steps)
+        if not isinstance(self.state, MathyAPIModelState):
+            raise ValueError(f"unknown state type: {type(self.state)}!")
+        return self.simplify_model(model=model, problem=problem, max_steps=max_steps)
+
+    def simplify_swarm(self, *, problem: str, max_steps: int) -> EpisodeMemory:
+        assert isinstance(self.state, MathyAPISwarmState), "not configured for swarm"
+        return swarm_solve(problem, self.state.config)
+
+    def simplify_model(
+        self, *, model: str = "mathy_alpha_sm", problem: str, max_steps: int,
     ) -> EpisodeMemory:
         """Simplify an input problem using the PolySimplify environment.
         
@@ -47,6 +75,7 @@ def simplify(
             to the solution for the input problem.
 
         """
+        assert isinstance(self.state, MathyAPIModelState), "not configured for model"
         import gym
         import tensorflow as tf
         from colr import color
@@ -66,12 +95,12 @@ def simplify(
         last_text = env.state.agent.problem
         last_action = -1
         last_reward = 0.0
-        selector = GreedyActionSelector(model=self.model, episode=0, worker_id=0)
+        selector = GreedyActionSelector(model=self.state.model, episode=0, worker_id=0)
         done = False
         while not done:
             env.render(last_action=last_action, last_reward=last_reward)
             window = episode_memory.to_window_observation(
-                last_observation, window_size=self.config.prediction_window_size
+                last_observation, window_size=self.state.config.prediction_window_size
             )
             try:
                 action, value = selector.select(
diff --git a/libraries/mathy_python/mathy/cli.py b/libraries/mathy_python/mathy/cli.py
@@ -29,6 +29,13 @@ def cli_contribute():
 
 
 @cli.command("simplify")
+@click.option(
+    "swarm",
+    "--swarm",
+    default=False,
+    is_flag=True,
+    help="Use swarm solver from fragile library without a trained model",
+)
 @click.option(
     "model", "--model", default="mathy_alpha_sm", help="The path to a mathy model",
 )
@@ -40,14 +47,23 @@ def cli_contribute():
     help="The max number of steps before the episode is over",
 )
 @click.argument("problem", type=str)
-def cli_simplify(agent: str, problem: str, model: str, max_steps: int):
+def cli_simplify(agent: str, problem: str, model: str, max_steps: int, swarm: bool):
     """Simplify an input polynomial expression."""
     setup_tf_env()
 
     from .models import load_model
     from .api import Mathy
+    from .agents.fragile import SwarmConfig
+
+    mt: Mathy
+    if swarm is True:
+        mt = Mathy(config=SwarmConfig())
+    else:
+        try:
+            mt = load_model(model)
+        except ValueError:
+            mt = Mathy(config=SwarmConfig())
 
-    mt: Mathy = load_model(model)
     mt.simplify(problem=problem, max_steps=max_steps)
 
 
diff --git a/libraries/mathy_python/mathy/envs/gym/mathy_gym_env.py b/libraries/mathy_python/mathy/envs/gym/mathy_gym_env.py
@@ -28,6 +28,8 @@ def __init__(
         self,
         env_class: Type[MathyEnv] = MathyEnv,
         env_problem_args: Optional[MathyEnvProblemArgs] = None,
+        env_problem: Optional[str] = None,
+        env_max_moves: int = 64,
         np_observation: bool = False,
         repeat_problem: bool = False,
         **env_kwargs,
@@ -38,7 +40,12 @@ def __init__(
         self.mathy = env_class(**env_kwargs)
         self.env_class = env_class
         self.env_problem_args = env_problem_args
-        self._challenge, _ = self.mathy.get_initial_state(env_problem_args)
+        if env_problem is not None:
+            self._challenge = MathyEnvState(
+                problem=env_problem, max_moves=env_max_moves
+            )
+        else:
+            self._challenge, _ = self.mathy.get_initial_state(env_problem_args)
         self.action_space = MaskedDiscrete(self.action_size, [1] * self.action_size)
 
     @property
diff --git a/libraries/mathy_python/requirements.txt b/libraries/mathy_python/requirements.txt
