def swarm_solve( problems: Union[List[str], str], config: SwarmConfig, max_steps: Union[List[int], int] = 128, silent: bool = False, ) -> Swarm: single_problem: bool = isinstance(problems, str) if single_problem: problems = [problems] if isinstance(max_steps, int): max_steps = [max_steps ] if single_problem else [max_steps] * len(problems) assert len(problems) > 0, "no problems to solve" assert len(problems) == len(max_steps) assert isinstance(problems, list) current_problem: str = problems.pop(0) current_max_moves: int = max_steps.pop(0) def env_callable(): nonlocal current_problem, current_max_moves return FragileMathyEnv( name="mathy_v0", problem=current_problem, repeat_problem=True, max_steps=current_max_moves, ) mathy_env: MathyEnv = env_callable()._env._env.mathy swarm: Swarm = mathy_swarm(config, env_callable) while True: if not silent: with msg.loading(f"Solving {current_problem} ..."): swarm.run() else: swarm.run() if not silent: if swarm.walkers.best_reward > EnvRewards.WIN: last_state = MathyEnvState.from_np( swarm.walkers.states.best_state) msg.good( f"Solved! {current_problem} = {last_state.agent.problem}") mathy_env.print_history(last_state) else: msg.fail(f"Failed to find a solution :(") if len(max_steps) > 0: current_max_moves = max_steps.pop(0) current_problem = problems.pop(0) else: break return swarm
def set_state(self, state: np.ndarray): assert self._env is not None, "env required to set_state" self._env.state = MathyEnvState.from_np(state) return state