fix(rewards): restore reentrant state scaling

 - turns out if you don't scale up the negative rewards for reentering states, you get a model that wins but only at the last moment after spinning in circles a bunch.
 - observe that associative swap tends to create trees where it's difficult for the model to commute terms. add negative reward for it

Author: justindujardin

libraries/mathy_python/mathy/env.py

@@ -115,7 +115,7 @@ def get_rewarding_actions(self, state: MathyEnvState) -> List[Type[BaseRule]]:
     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 []
+        return [AssociativeSwapRule]
 
     def max_moves_fn(
         self, problem: MathyEnvProblem, config: MathyEnvProblemArgs
@@ -204,8 +204,11 @@ def get_state_transition(
             if list_count <= 1 or key != expression.raw:
                 continue
 
+            # NOTE: the reward is scaled by how many times this state has been visited
             return time_step.transition(
-                features, reward=EnvRewards.PREVIOUS_LOCATION, discount=self.discount,
+                features,
+                reward=EnvRewards.PREVIOUS_LOCATION * list_count,
+                discount=self.discount,
             )
 
         if len(agent.history) > 0: