Spaces:
Running
Running
File size: 1,594 Bytes
e11e4fe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
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
|