feat: Add support for integrating with Fragile library

* chore: stash WIP fragile search mathy agent

 - add `mathy_fragile.py` in the root for now. It requires a local mathy installation to use the hacked up library compatible with plangym: https://mathy.ai/contributing/#use-your-local-version-of-mathy
 - requires plangym/fragile packages, installed by whatever means you use

* chore(env): add invalid move error scaling

 - calculate the distance from the selected invalid action to the nearest valid action and provide a negative reward based on that.
  - remove some extra state from the MathyGymEnv to ensure that fragile can cleanly reset all of the state during batch steps

* chore: finish cleaning up gym env for clean reset

 - I changed the env to reset to the same initial state when called. I think this is more in-line with what an atari env might do.

* chore: use node types as observations

 - sequence of node type IDs
 - adjust scale of rewards to remove negatives

* test(state): fix to_string/from_string

* refactor(env): use positive reward signals for fragile search

 - normalize all rewards so that dead-states return 0.0 and positive ones return > 0.0 (as in paper)
 - treat invalid action selections as null states and continue the simulation

* chore: cleanup and fix test script

* chore: cleanup

* chore(fragile): add masked action selection

 - this is pretty slow by comparison. See if there's a way to avoid looping over the action masks. Maybe build the probabilities into the observation so the calculation is unnecessary?

* chore: faster numpy action selection

- calculate mask probabilities and return as observation
- use numpy trickery to avoid looping over the observations to choose random actions from probabilities in batches

* chore: fix fragile oob/terminal conditions

* chore: fix tests

* chore: drop fragile script

Author: justindujardin

libraries/mathy_pydoc/requirements-dev.txt

@@ -0,0 +1,5 @@
+mypy
+pytest
+pytest-cov
+black
+flake8

libraries/mathy_pydoc/tools/setup.sh

@@ -10,6 +10,6 @@ fi
 
 . .env/bin/activate
 echo "Installing/updating requirements..."
-pip install pytest pytest-cov
-pip install -e .[dev]
+pip install -e .
+pip install -r requirements-dev.txt
 

libraries/mathy_python/mathy/agents/a3c/agent.py

@@ -10,6 +10,8 @@
 from ..policy_value_model import get_or_create_policy_model, PolicyValueModel
 from .config import A3CConfig
 from .worker import A3CWorker
+from ...envs.gym import MathyGymEnv
+from ...state import observations_to_window
 
 
 class A3CAgent:
@@ -33,17 +35,17 @@ def __init__(self, args: A3CConfig, env_extra: dict = None):
             win_threshold=self.args.teacher_promote_wins,
             lose_threshold=self.args.teacher_demote_wins,
         )
-        env = gym.make(self.teacher.get_env(0, 0), **self.env_extra)
+        env: MathyGymEnv = gym.make(self.teacher.get_env(0, 0), **self.env_extra)
         self.action_size = env.action_space.n
         self.log_dir = os.path.join(self.args.model_dir, "tensorboard")
         self.writer = tf.summary.create_file_writer(self.log_dir)
-        init_window = env.initial_window()
+        initial_window = observations_to_window([env.reset()])
         self.global_model = get_or_create_policy_model(
             args=args, predictions=self.action_size, is_main=True, env=env.mathy
         )
         with self.writer.as_default():
             tf.summary.trace_on(graph=True)
-            inputs = init_window.to_inputs()
+            inputs = initial_window.to_inputs()
 
             @tf.function
             def trace_fn():

libraries/mathy_python/mathy/agents/a3c/worker.py

@@ -224,13 +224,13 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
         last_observation: MathyObservation = env.reset()
         last_text = env.state.agent.problem
         last_action = -1
-        last_reward = -1
+        last_reward = 0.0
 
         selector = self.build_episode_selector(env)
 
         while not done and A3CWorker.request_quit is False:
             if self.args.print_training and self.worker_idx == 0:
-                env.render(self.args.print_mode, None)
+                env.render(last_action=last_action, last_reward=last_reward)
             window = episode_memory.to_window_observation(
                 last_observation, window_size=self.args.prediction_window_size
             )
@@ -253,7 +253,7 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
             )
             if time_count == self.args.update_gradients_every or done:
                 if done and self.args.print_training and self.worker_idx == 0:
-                    env.render(self.args.print_mode, None)
+                    env.render(last_action=last_action, last_reward=last_reward)
 
                 self.update_global_network(done, observation, episode_memory)
                 self.maybe_write_histograms()

libraries/mathy_python/mathy/api.py

@@ -65,11 +65,11 @@ def simplify(
         assert env.state is not None
         last_text = env.state.agent.problem
         last_action = -1
-        last_reward = -1
+        last_reward = 0.0
         selector = GreedyActionSelector(model=self.model, episode=0, worker_id=0)
         done = False
         while not done:
-            env.render(self.config.print_mode, None)
+            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
             )
@@ -96,7 +96,7 @@ def simplify(
             if done:
                 # Last timestep reward
                 win = reward > 0.0
-                env.render(self.config.print_mode, None)
+                env.render(last_action=last_action, last_reward=last_reward)
                 print(
                     color(
                         text="SOLVE" if win else "FAIL", fore="green" if win else "red",

libraries/mathy_python/mathy/core/rule.py

@@ -83,12 +83,13 @@ class ExpressionChangeRule:
 
     rule: BaseRule
     node: Optional[MathExpression]
-    result: MathExpression
+    result: Optional[MathExpression]
     _save_parent: Optional[MathExpression]
 
     def __init__(self, rule, node: MathExpression = None):
         self.rule = rule
         self.node = node
+        self.result = None
         self._save_parent = None
 
     def save_parent(

libraries/mathy_python/mathy/env.py

@@ -43,11 +43,13 @@ def __init__(
         rules: List[BaseRule] = None,
         max_moves: int = 20,
         verbose: bool = False,
+        error_invalid: bool = False,
         reward_discount: float = 0.99,
     ):
         self.discount = reward_discount
         self.verbose = verbose
         self.max_moves = max_moves
+        self.error_invalid = error_invalid
         self.parser = ExpressionParser()
         if rules is None:
             self.rules = MathyEnv.core_rules()
@@ -73,18 +75,6 @@ def action_size(self) -> int:
         """Return the number of available actions"""
         return len(self.rules)
 
-    def step(
-        self, state: MathyEnvState, action: int, as_observation: bool = False
-    ) -> Tuple[Union[MathyEnvState, MathyObservation], float, bool, Any]:
-        new_state, transition, change = self.get_next_state(state, action)
-        observation = self.state_to_observation(state)
-        info = {"transition": transition}
-        done = is_terminal_transition(transition)
-        self.last_action = action
-        self.last_change = change
-        self.last_reward = round(float(transition.reward), 4)
-        return observation, transition.reward, done, info
-
     def finalize_state(self, state: MathyEnvState):
         """Perform final checks on a problem state, to ensure the episode yielded
         results that were uncorrupted by transformation errors."""
@@ -108,15 +98,17 @@ def get_rewarding_actions(self, state: MathyEnvState) -> List[Type[BaseRule]]:
         #       rewarding.
         return [
             ConstantsSimplifyRule,
-            DistributiveMultiplyRule,
             DistributiveFactorOutRule,
             VariableMultiplyRule,
         ]
 
     def get_penalizing_actions(self, state: MathyEnvState) -> List[Type[BaseRule]]:
         """Get the list of penalizing action types. When these actions
         are selected, the agent gets a negative reward."""
-        return [AssociativeSwapRule]
+        return [
+            AssociativeSwapRule,
+            DistributiveMultiplyRule,
+        ]
 
     def max_moves_fn(
         self, problem: MathyEnvProblem, config: MathyEnvProblemArgs
@@ -267,28 +259,41 @@ def get_next_state(
         op_not_rule = not isinstance(operation, BaseRule)
         op_cannot_apply = token is None or operation.can_apply_to(token) is False
         if token is None or op_not_rule or op_cannot_apply:
-            steps = int(env_state.max_moves - agent.moves_remaining)
-            msg = "Step: {} - Invalid action({}) '{}' for expression '{}'.".format(
-                steps, action, type(operation), expression
-            )
-            raise_with_history("Invalid Action", msg, agent.history)
-            raise ValueError(f"Invalid Action: {msg}")
+            if self.error_invalid:
+                steps = int(env_state.max_moves - agent.moves_remaining)
+                msg = "Step: {} - Invalid action({}) '{}' for expression '{}'.".format(
+                    steps, action, type(operation), expression
+                )
+                raise_with_history("Invalid Action", msg, agent.history)
+                raise ValueError(f"Invalid Action: {msg}")
+            else:
+                valid_mask = self.get_valid_moves(env_state)
+                # Non-masked searches ignore invalid moves entirely
+                out_env = MathyEnvState.copy(env_state)
+                out_env.action = -1
+                obs = out_env.to_observation(
+                    self.get_valid_moves(out_env), parser=self.parser
+                )
+                transition = time_step.transition(obs, EnvRewards.INVALID_MOVE)
+                return out_env, transition, ExpressionChangeRule(BaseRule())
 
         change = operation.apply_to(token.clone_from_root())
         root = change.result.get_root()
         change_name = operation.name
         out_problem = str(root)
         out_env = env_state.get_out_state(
             problem=out_problem,
-            focus_index=token_index,
+            focus=token_index,
             action=action_index,
             moves_remaining=agent.moves_remaining - 1,
         )
 
+        transition = self.get_state_transition(out_env, searching)
         if not searching and self.verbose:
             token_idx = int("{}".format(token_index).zfill(3))
-            self.print_state(out_env, change_name[:25].lower(), token_idx, change)
-        transition = self.get_state_transition(out_env, searching)
+            self.print_state(
+                out_env, change_name[:25].lower(), token_idx, change, transition.reward
+            )
         return out_env, transition, change
 
     def print_state(
@@ -316,7 +321,7 @@ def render_state(
     ):
         """Render the given state to a string suitable for printing to a log"""
         changed_problem = env_state.agent.problem
-        if change is not None:
+        if change is not None and change.result is not None:
             changed_problem = change.result.get_root().terminal_text
         output = """{:<25} | {}""".format(action_name.lower(), changed_problem)
 
@@ -382,16 +387,16 @@ def get_agent_actions_count(self, env_state: MathyEnvState) -> int:
         return self.action_size * node_count
 
     def get_token_at_index(
-        self, expression: MathExpression, focus_index: int
+        self, expression: MathExpression, index: int
     ) -> Optional[MathExpression]:
-        """Get the token that is `focus_index` from the left of the expression"""
+        """Get the token that is `index` from the left of the expression"""
         count = 0
         result = None
 
         def visit_fn(node, depth, data):
             nonlocal result, count
             result = node
-            if count == focus_index:
+            if count == index:
                 return STOP
             count = count + 1