您最多选择25个主题
主题必须以中文或者字母或数字开头，可以包含连字符 (-)，并且长度不得超过35个字符
89 行
3.4 KiB
89 行
3.4 KiB
from typing import Dict, List
|
|
|
|
from mlagents_envs.base_env import BaseEnv
|
|
from mlagents.trainers.env_manager import EnvManager, EnvironmentStep
|
|
from mlagents_envs.timers import timed
|
|
from mlagents.trainers.action_info import ActionInfo
|
|
from mlagents.trainers.brain import BrainParameters, AllBrainInfo
|
|
from mlagents_envs.side_channel.float_properties_channel import FloatPropertiesChannel
|
|
from mlagents.trainers.brain_conversion_utils import (
|
|
step_result_to_brain_info,
|
|
group_spec_to_brain_parameters,
|
|
)
|
|
|
|
|
|
class SimpleEnvManager(EnvManager):
|
|
"""
|
|
Simple implementation of the EnvManager interface that only handles one BaseEnv at a time.
|
|
This is generally only useful for testing; see SubprocessEnvManager for a production-quality implementation.
|
|
"""
|
|
|
|
def __init__(self, env: BaseEnv, float_prop_channel: FloatPropertiesChannel):
|
|
super().__init__()
|
|
self.shared_float_properties = float_prop_channel
|
|
self.env = env
|
|
self.previous_step: EnvironmentStep = EnvironmentStep({}, {}, {})
|
|
self.previous_all_action_info: Dict[str, ActionInfo] = {}
|
|
|
|
def step(self) -> List[EnvironmentStep]:
|
|
all_action_info = self._take_step(self.previous_step)
|
|
self.previous_all_action_info = all_action_info
|
|
|
|
for brain_name, action_info in all_action_info.items():
|
|
self.env.set_actions(brain_name, action_info.action)
|
|
self.env.step()
|
|
all_brain_info = self._generate_all_brain_info()
|
|
step_brain_info = all_brain_info
|
|
|
|
step_info = EnvironmentStep(
|
|
self.previous_step.current_all_brain_info,
|
|
step_brain_info,
|
|
self.previous_all_action_info,
|
|
)
|
|
self.previous_step = step_info
|
|
return [step_info]
|
|
|
|
def reset(
|
|
self, config: Dict[str, float] = None
|
|
) -> List[EnvironmentStep]: # type: ignore
|
|
if config is not None:
|
|
for k, v in config.items():
|
|
self.shared_float_properties.set_property(k, v)
|
|
self.env.reset()
|
|
all_brain_info = self._generate_all_brain_info()
|
|
self.previous_step = EnvironmentStep({}, all_brain_info, {})
|
|
return [self.previous_step]
|
|
|
|
@property
|
|
def external_brains(self) -> Dict[str, BrainParameters]:
|
|
result = {}
|
|
for brain_name in self.env.get_agent_groups():
|
|
result[brain_name] = group_spec_to_brain_parameters(
|
|
brain_name, self.env.get_agent_group_spec(brain_name)
|
|
)
|
|
return result
|
|
|
|
@property
|
|
def get_properties(self) -> Dict[str, float]:
|
|
return self.shared_float_properties.get_property_dict_copy()
|
|
|
|
def close(self):
|
|
self.env.close()
|
|
|
|
@timed
|
|
def _take_step(self, last_step: EnvironmentStep) -> Dict[str, ActionInfo]:
|
|
all_action_info: Dict[str, ActionInfo] = {}
|
|
for brain_name, brain_info in last_step.current_all_brain_info.items():
|
|
all_action_info[brain_name] = self.policies[brain_name].get_action(
|
|
brain_info
|
|
)
|
|
return all_action_info
|
|
|
|
def _generate_all_brain_info(self) -> AllBrainInfo:
|
|
all_brain_info = {}
|
|
for brain_name in self.env.get_agent_groups():
|
|
all_brain_info[brain_name] = step_result_to_brain_info(
|
|
self.env.get_step_result(brain_name),
|
|
self.env.get_agent_group_spec(brain_name),
|
|
)
|
|
return all_brain_info
|