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
