Spaces:
Sleeping
Sleeping
from typing import Optional, Dict, Any, Tuple | |
from gym import error | |
from mlagents_envs.base_env import BaseEnv | |
from pettingzoo import ParallelEnv | |
from mlagents_envs.envs.unity_pettingzoo_base_env import UnityPettingzooBaseEnv | |
class UnityParallelEnv(UnityPettingzooBaseEnv, ParallelEnv): | |
""" | |
Unity Parallel (PettingZoo) environment wrapper. | |
""" | |
def __init__(self, env: BaseEnv, seed: Optional[int] = None): | |
""" | |
Initializes a Unity Parallel environment wrapper. | |
:param env: The UnityEnvironment that is being wrapped. | |
:param seed: The seed for the action spaces of the agents. | |
""" | |
super().__init__(env, seed) | |
def reset(self) -> Dict[str, Any]: | |
""" | |
Resets the environment. | |
""" | |
super().reset() | |
return self._observations | |
def step(self, actions: Dict[str, Any]) -> Tuple: | |
self._assert_loaded() | |
if len(self._live_agents) <= 0 and actions: | |
raise error.Error( | |
"You must reset the environment before you can perform a step." | |
) | |
# Process actions | |
for current_agent, action in actions.items(): | |
self._process_action(current_agent, action) | |
# Reset reward | |
for k in self._rewards.keys(): | |
self._rewards[k] = 0 | |
# Step environment | |
self._step() | |
# Agent cleanup and sorting | |
self._cleanup_agents() | |
self._live_agents.sort() # unnecessary, only for passing API test | |
return self._observations, self._rewards, self._dones, self._infos | |