MultiAgent

class MultiAgent:
    def __init__(self, agents: List[Agent], environment: Environment)

@property
def step_count(self) -> int

def step(self) -> Dict[str, Any]

def reset(self) -> Dict[str, Any]

def add_agent(self, agent: Agent) -> None

def remove_agent(self, agent_id: str) -> Optional[Agent]

def __len__(self) -> int

import neurenix as nx
from neurenix.agent import Agent, MultiAgent, Environment

# Define a custom environment
class GridWorld(Environment):
    def __init__(self, width, height):
        super().__init__()
        self.width = width
        self.height = height
        self.agent_positions = {}
    
    def reset(self):
        # Reset environment state
        self.agent_positions = {}
        return {}
    
    def observe(self, agent):
        # Return observation for a specific agent
        position = self.agent_positions.get(agent.id, (0, 0))
        return nx.Tensor([position[0], position[1]])
    
    def step(self, actions):
        # Apply all agent actions
        rewards = {}
        for agent_id, action in actions.items():
            # Update agent position based on action
            # Calculate rewards
            rewards[agent_id] = self.calculate_reward(agent_id, action)
        
        return {
            "rewards": rewards,
            "done": self.is_terminal(),
            "info": {}
        }

# Create agents
agent1 = MyAgent(name="Agent-1")
agent2 = MyAgent(name="Agent-2")
agent3 = MyAgent(name="Agent-3")

# Create environment
env = GridWorld(width=10, height=10)

# Create multi-agent system
mas = MultiAgent(
    agents=[agent1, agent2, agent3],
    environment=env
)

print(f"Number of agents: {len(mas)}")

# Run simulation
obs = mas.reset()

for episode in range(100):
    obs = mas.reset()
    done = False
    episode_rewards = {agent.id: 0 for agent in mas.agents}
    
    while not done:
        # Step the multi-agent system
        results = mas.step()
        
        # Track rewards
        for agent_id, reward in results["rewards"].items():
            episode_rewards[agent_id] += reward
        
        # Check if done
        done = results["done"]
        
        # Train agents
        for agent in mas.agents:
            agent.learn(results)
    
    print(f"Episode {episode}: {episode_rewards}")

# Add a new agent dynamically
new_agent = MyAgent(name="Agent-4")
mas.add_agent(new_agent)
print(f"Agents after addition: {len(mas)}")

# Remove an agent
removed = mas.remove_agent("Agent-1")
print(f"Removed: {removed.name}")
print(f"Remaining agents: {len(mas)}")