Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import json | |
| from dataclasses import dataclass, field | |
| from typing import Any, Dict, Optional, Tuple | |
| from poke_env.environment.battle import Battle | |
| from poke_env.player.player import Player | |
| from .action_space import ( | |
| ActionJSON, | |
| ActionOption, | |
| build_action_instructions, | |
| enumerate_actions, | |
| extract_action_json_from_text, | |
| parse_llm_action, | |
| ) | |
| from .config import EnvConfig | |
| from .pokeenv_client import PokeEnvClient | |
| from .reward import ( | |
| BattleStateSummary, | |
| ILLEGAL_ACTION_PENALTY, | |
| RewardTrackingState, | |
| calculate_reward, | |
| count_new_passive_hits_for_turn, | |
| summarize_battle_state, | |
| ) | |
| from .state_formatter import OpponentHistoryTracker, format_battle_state | |
| class PokemonShowdownEnv: | |
| """Synchronous, OpenEnv-style wrapper around a poke-env battle. | |
| The environment exposes a simple Gymnasium-like / OpenEnv-like API: | |
| obs = env.reset() | |
| obs, reward, done, info = env.step(action_json_str) | |
| where `action_json_str` is a JSON string describing a move or switch using | |
| the constrained 9-action space. | |
| """ | |
| config: EnvConfig = field(default_factory=EnvConfig) | |
| _client: PokeEnvClient = field(init=False) | |
| _opponent_history: OpponentHistoryTracker = field(init=False) | |
| _reward_trackers: RewardTrackingState = field(init=False) | |
| _prev_state: Optional[BattleStateSummary] = field(init=False, default=None) | |
| _steps_this_battle: int = field(init=False, default=0) | |
| # Running total of passive hits — updated O(k) per step via the single-turn | |
| # scanner, never by re-scanning the full observation history. | |
| _cumulative_passive_hits: int = field(init=False, default=0) | |
| _battle_index: int = field(init=False, default=0) | |
| _battle_reward_total: float = field(init=False, default=0.0) | |
| _no_progress_steps: int = field(init=False, default=0) | |
| def __post_init__(self) -> None: | |
| self._client = PokeEnvClient(config=self.config) | |
| self._opponent_history = OpponentHistoryTracker() | |
| self._reward_trackers = RewardTrackingState() | |
| def _log(self, message: str) -> None: | |
| if self.config.verbose_logging: | |
| print(f"[PokemonShowdownEnv] {message}", flush=True) | |
| # ------------------------------------------------------------------ API | |
| def reset(self) -> str: | |
| """Start a new battle and return the initial markdown state.""" | |
| self._battle_index += 1 | |
| self._client.start_new_battle() | |
| self._opponent_history = OpponentHistoryTracker() | |
| self._reward_trackers = RewardTrackingState() | |
| self._steps_this_battle = 0 | |
| self._cumulative_passive_hits = 0 | |
| self._battle_reward_total = 0.0 | |
| self._no_progress_steps = 0 | |
| battle = self._wait_for_battle_or_raise() | |
| self._log( | |
| f"Battle {self._battle_index} started at turn={battle.turn} " | |
| f"(format={self.config.battle_format})." | |
| ) | |
| self._prev_state = summarize_battle_state(battle, self._cumulative_passive_hits) | |
| return format_battle_state(battle, self._opponent_history) | |
| def step(self, action_json: str | Dict[str, Any]) -> Tuple[str, float, bool, Dict[str, Any]]: | |
| """Apply one action and return (state_str, reward, done, info).""" | |
| battle = self._ensure_battle() | |
| if battle.finished: | |
| raise RuntimeError("Cannot call step() on a finished battle. Call reset().") | |
| self._steps_this_battle += 1 | |
| if self._steps_this_battle > self.config.max_steps_per_battle: | |
| return self._terminal_from_truncation(battle) | |
| valid_actions = enumerate_actions(battle) | |
| if isinstance(action_json, dict): | |
| raw = json.dumps(action_json) | |
| else: | |
| raw = action_json | |
| used_fallback = False | |
| try: | |
| parsed = parse_llm_action(raw, valid_actions) | |
| order = self._to_battle_order(parsed, valid_actions, battle) | |
| except ValueError: | |
| extracted = extract_action_json_from_text(raw) | |
| if extracted is not None: | |
| try: | |
| parsed = parse_llm_action(extracted, valid_actions) | |
| order = self._to_battle_order(parsed, valid_actions, battle) | |
| except ValueError: | |
| used_fallback = True | |
| else: | |
| used_fallback = True | |
| if used_fallback: | |
| opt = valid_actions[0] | |
| from poke_env.player import Player as PlayerCls | |
| if opt.action_type == "move" and opt.move is not None: | |
| order = PlayerCls.create_order(opt.move) | |
| else: | |
| order = PlayerCls.create_order(opt.pokemon) | |
| previous_turn = battle.turn | |
| self._client.send_action(order) | |
| new_battle = self._client.wait_for_battle_update(previous_turn) or battle | |
| # Increment the passive-hit counter by scanning only the turn that just | |
| # resolved — O(k) where k = events on that single turn, not O(total turns). | |
| self._cumulative_passive_hits += count_new_passive_hits_for_turn( | |
| new_battle, previous_turn | |
| ) | |
| prev_state = self._prev_state or summarize_battle_state(battle, self._cumulative_passive_hits) | |
| curr_state = summarize_battle_state(new_battle, self._cumulative_passive_hits) | |
| active = new_battle.active_pokemon | |
| opponent_active = new_battle.opponent_active_pokemon | |
| if used_fallback: | |
| reward = ILLEGAL_ACTION_PENALTY | |
| else: | |
| reward = calculate_reward( | |
| prev_state=prev_state, | |
| curr_state=curr_state, | |
| action=ActionJSON(action=parsed.action, choice=parsed.choice), | |
| trackers=self._reward_trackers, | |
| active=active, | |
| opponent_active=opponent_active, | |
| ) | |
| # Small time cost per turn to discourage excessively long battles. | |
| reward += self.config.step_living_penalty | |
| self._prev_state = curr_state | |
| if new_battle.turn == previous_turn and not new_battle.finished: | |
| self._no_progress_steps += 1 | |
| else: | |
| self._no_progress_steps = 0 | |
| done_reason: Optional[str] = None | |
| done = False | |
| if new_battle.finished: | |
| done = True | |
| done_reason = "battle_finished" | |
| elif self._steps_this_battle >= self.config.max_steps_per_battle: | |
| done = True | |
| done_reason = "max_steps" | |
| reward += self.config.max_steps_termination_penalty | |
| elif (self._battle_reward_total + reward) <= self.config.min_battle_reward: | |
| done = True | |
| done_reason = "min_battle_reward" | |
| elif self._no_progress_steps >= self.config.max_no_progress_steps: | |
| done = True | |
| done_reason = "no_progress_timeout" | |
| reward += self.config.no_progress_termination_penalty | |
| self._battle_reward_total += reward | |
| # If we terminate early (not a natural finished battle), forfeit cleanly | |
| # so the next reset starts from a free player/session state. | |
| if done and not new_battle.finished and done_reason in { | |
| "max_steps", | |
| "min_battle_reward", | |
| "no_progress_timeout", | |
| }: | |
| try: | |
| self._client.forfeit_current_battle() | |
| except Exception: | |
| pass | |
| obs = format_battle_state(new_battle, self._opponent_history) | |
| info: Dict[str, Any] = { | |
| "turn": new_battle.turn, | |
| "valid_actions": [ | |
| {"action": a.action_type, "choice": a.choice} for a in valid_actions | |
| ], | |
| "instructions": build_action_instructions(valid_actions), | |
| "battle_finished": new_battle.finished, | |
| "reason": done_reason, | |
| "action_illegal": used_fallback, | |
| "battle_reward_total": self._battle_reward_total, | |
| "no_progress_steps": self._no_progress_steps, | |
| } | |
| if self.config.verbose_logging: | |
| should_log_step = ( | |
| used_fallback | |
| or done | |
| or self._steps_this_battle == 1 | |
| or self._steps_this_battle % max(1, self.config.log_every_n_steps) == 0 | |
| ) | |
| if should_log_step: | |
| self._log( | |
| f"battle={self._battle_index} step={self._steps_this_battle} " | |
| f"turn={new_battle.turn} reward={reward:.3f} " | |
| f"running_reward={self._battle_reward_total:.3f} " | |
| f"illegal_action={used_fallback} done={done}" | |
| ) | |
| return obs, reward, done, info | |
| # ------------------------------------------------------------------ helpers | |
| def _wait_for_battle_or_raise(self) -> Battle: | |
| battle = self._client.battle | |
| if battle is None: | |
| battle = self._client.wait_for_battle_update(previous_turn=0) | |
| if battle is None: | |
| raise RuntimeError("Failed to obtain initial battle from poke-env.") | |
| return battle | |
| def _ensure_battle(self) -> Battle: | |
| battle = self._client.battle | |
| if battle is None: | |
| raise RuntimeError("No active battle. Call reset() first.") | |
| return battle | |
| def _terminal_from_truncation(self, battle: Battle) -> Tuple[str, float, bool, Dict[str, Any]]: | |
| obs = format_battle_state(battle, self._opponent_history) | |
| info: Dict[str, Any] = { | |
| "turn": battle.turn, | |
| "battle_finished": battle.finished, | |
| "reason": "max_steps", | |
| } | |
| return obs, self.config.max_steps_termination_penalty, True, info | |
| def _to_battle_order( | |
| parsed: ActionJSON, | |
| valid_actions: list[ActionOption], | |
| battle: Battle, | |
| ) -> "Player.create_order.__annotations__['return']": | |
| from poke_env.player import Player as PlayerCls | |
| for opt in valid_actions: | |
| if opt.action_type == parsed.action and opt.choice == parsed.choice: | |
| if opt.action_type == "move" and opt.move is not None: | |
| return PlayerCls.create_order(opt.move) | |
| if opt.action_type == "switch" and opt.pokemon is not None: | |
| return PlayerCls.create_order(opt.pokemon) | |
| raise ValueError(f"Could not map parsed action {parsed} to a BattleOrder") | |