Spaces:
Build error
Build error
| from __future__ import annotations | |
| from typing import TYPE_CHECKING, List | |
| import numpy as np | |
| import json | |
| from agentverse.message import Message | |
| from . import selector_registry as SelectorRegistry | |
| from .base import BaseSelector | |
| if TYPE_CHECKING: | |
| from agentverse.environments import PokemonEnvironment | |
| class PokemonSelector(BaseSelector): | |
| """ | |
| Selector for Pokemon environment | |
| """ | |
| def select_message( | |
| self, environment: PokemonEnvironment, messages: List[Message] | |
| ) -> List[Message]: | |
| valid = [] | |
| talk_matrix = np.zeros((len(environment.agents), len(environment.agents))) | |
| agent_to_idx = {agent.name: i for i, agent in enumerate(environment.agents)} | |
| for i, message in enumerate(messages): | |
| try: | |
| content = json.loads(message.content) | |
| except json.decoder.JSONDecodeError: | |
| valid.append(0) | |
| continue | |
| if content["action"] == "Speak": | |
| try: | |
| if "to" not in content: | |
| # If the model does not generate receiver, then we discard the message | |
| valid.append(0) | |
| elif content["to"] in agent_to_idx: | |
| # TODO: allow talk to a list of agents | |
| valid.append(1) | |
| # talk_matrix[i][j] = 1 ==> i talk to j | |
| talk_matrix[agent_to_idx[message.sender]][ | |
| agent_to_idx[content["to"]] | |
| ] = 1 | |
| else: | |
| # If the receiver is not in the environment, then we discard the message | |
| valid.append(0) | |
| except: | |
| valid.append(0) | |
| continue | |
| elif content["action"] == "MoveTo": | |
| # If the agent move to a location that does not exist, then we discard the message | |
| valid.append( | |
| "to" in content and content["to"] in environment.locations_to_agents | |
| ) | |
| else: | |
| valid.append(1) | |
| selected_messages = [] | |
| for i, message in enumerate(messages): | |
| content = json.loads(message.content) | |
| sender_idx = agent_to_idx[message.sender] | |
| if valid[i] == 0: | |
| selected_messages.append(Message()) | |
| continue | |
| if content["action"] == "MoveTo": | |
| if np.sum(talk_matrix[:, sender_idx]) > 0: | |
| # If someone talk to this agent, then we discard the move action | |
| selected_messages.append(Message()) | |
| else: | |
| selected_messages.append(message) | |
| elif content["action"] == "Speak": | |
| receiver_idx = agent_to_idx[content["to"]] | |
| if talk_matrix[sender_idx][receiver_idx] == 0: | |
| # If this agent talk to someone who also talk to this agent, and we | |
| # select the message from this agent, then we discard the message | |
| selected_messages.append(Message()) | |
| continue | |
| if np.sum(talk_matrix[receiver_idx, :]) > 0: | |
| if talk_matrix[receiver_idx][sender_idx] == 1: | |
| # If the receiver talk to this agent, then we randomly select one message | |
| if sender_idx < receiver_idx: | |
| if np.random.random() < 0.5: | |
| selected_messages.append(message) | |
| talk_matrix[receiver_idx][sender_idx] = 0 | |
| else: | |
| selected_messages.append(Message()) | |
| talk_matrix[sender_idx][receiver_idx] = 0 | |
| else: | |
| print("Shouldn't happen") | |
| else: | |
| # If the receiver talk to other agent, we still talk to the receiver (?) | |
| selected_messages.append(message) | |
| else: | |
| selected_messages.append(message) | |
| else: | |
| selected_messages.append(message) | |
| return selected_messages | |