feature(zms): add new league middlewares and other models and tools.

ding/envs/env_manager/env_supervisor.py

@@ -61,6 +61,7 @@ def __init__(
             episode_num: int = float("inf"),
             shared_memory: bool = True,
             copy_on_get: bool = True,
+            return_original_data: bool = False,
             **kwargs
     ) -> None:
         """
@@ -78,6 +79,9 @@ def __init__(
             - retry_waiting_time (:obj:`Optional[float]`): Wait time on each retry.
             - shared_memory (:obj:`bool`): Use shared memory in multiprocessing.
             - copy_on_get (:obj:`bool`): Use copy on get in multiprocessing.
+            - return_original_data (:obj:`bool`): Return original observation, 
+                so that the attribute self._ready_obs is not a tnp.array but only the original observation,
+                and the property self.ready_obs is a dict in which the key is the env_id.
         """
         if kwargs:
             logging.warning("Unknown parameters on env supervisor: {}".format(kwargs))
@@ -122,6 +126,7 @@ def __init__(
         self._retry_waiting_time = retry_waiting_time
         self._env_replay_path = None
         self._episode_num = episode_num
+        self._return_original_data = return_original_data
         self._init_states()
 
     def _init_states(self):
@@ -255,6 +260,8 @@ def ready_obs(self) -> tnp.array:
             >>> timesteps = env_manager.step(action)
         """
         active_env = [i for i, s in self._env_states.items() if s == EnvState.RUN]
+        if self._return_original_data:
+            return {i: self._ready_obs[i] for i in active_env}
         active_env.sort()
         obs = [self._ready_obs.get(i) for i in active_env]
         if len(obs) == 0:
@@ -409,16 +416,17 @@ def _recv_step_callback(
                         remain_payloads[p.req_id] = p
             # make the type and content of key as similar as identifier,
             # in order to call them as attribute (e.g. timestep.xxx), such as ``TimeLimit.truncated`` in cartpole info
-            info = make_key_as_identifier(info)
-            payload.data = tnp.array(
-                {
-                    'obs': obs,
-                    'reward': reward,
-                    'done': done,
-                    'info': info,
-                    'env_id': payload.proc_id
-                }
-            )
+            if not self._return_original_data:
+                info = make_key_as_identifier(info)
+                payload.data = tnp.array(
+                    {
+                        'obs': obs,
+                        'reward': reward,
+                        'done': done,
+                        'info': info,
+                        'env_id': payload.proc_id
+                    }
+                )
             self._ready_obs[payload.proc_id] = obs
         return payload
 

ding/framework/__init__.py

@@ -1,5 +1,6 @@
-from .context import Context, OnlineRLContext, OfflineRLContext
+from .context import Context, OnlineRLContext, OfflineRLContext, BattleContext
 from .task import Task, task
 from .parallel import Parallel
 from .event_loop import EventLoop
+from .event_enum import EventEnum
 from .supervisor import Supervisor

ding/framework/context.py

@@ -75,3 +75,38 @@ def __init__(self, *args, **kwargs) -> None:
         self.last_eval_iter = -1
 
         self.keep('train_iter', 'last_eval_iter')
+
+
+class BattleContext(Context):
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.__dict__ = self
+        # collect target paras
+        self.n_episode = None
+
+        #collect process paras
+        self.env_episode = 0
+        self.env_step = 0
+        self.total_envstep_count = 0
+        self.train_iter = 0
+        self.collect_kwargs = {}
+        self.current_policies = []
+
+        #job paras
+        self.player_id_list = []
+        self.job_finish = False
+
+        #data
+        self.obs = None
+        self.actions = None
+        self.inference_output = {}
+        self.trajectories = None
+
+        #Return data paras
+        self.episodes = []
+        self.episode_info = []
+        self.trajectories_list = []
+        self.train_data = None
+
+        self.keep('train_iter')

ding/framework/event_enum.py

@@ -0,0 +1,16 @@
+from enum import Enum, unique
+
+
+@unique
+class EventEnum(str, Enum):
+    # events emited by coordinators
+    COORDINATOR_DISPATCH_ACTOR_JOB = "on_coordinator_dispatch_actor_job_{actor_id}"
+
+    # events emited by learners
+    LEARNER_SEND_MODEL = "on_learner_send_model"
+    LEARNER_SEND_META = "on_learner_send_meta"
+
+    # events emited by actors
+    ACTOR_GREETING = "on_actor_greeting"
+    ACTOR_SEND_DATA = "on_actor_send_meta_player_{player}"
+    ACTOR_FINISH_JOB = "on_actor_finish_job"

ding/framework/event_loop.py

@@ -1,3 +1,4 @@
+import re
 from collections import defaultdict
 from typing import Callable, Optional
 from concurrent.futures import ThreadPoolExecutor
@@ -23,6 +24,12 @@ def on(self, event: str, fn: Callable) -> None:
             - event (:obj:`str`): Event name.
             - fn (:obj:`Callable`): The function.
         """
+        # check if the event name contains unfilled parameters.
+        params = re.findall(r"\{(.*?)\}", event)
+        if params:
+            raise ValueError(
+                "Event name missing parameters: {}. Please use String.format() to fill up".format(", ".join(params))
+            )
         self._listeners[event].append(fn)
 
     def off(self, event: str, fn: Optional[Callable] = None) -> None:
@@ -65,6 +72,12 @@ def emit(self, event: str, *args, **kwargs) -> None:
         """
         if self._exception:
             raise self._exception
+        # check if the event name contains unfilled parameters.
+        params = re.findall(r"\{(.*?)\}", event)
+        if params:
+            raise ValueError(
+                "Event name missing parameters: {}. Please use String.format() to fill up".format(", ".join(params))
+            )
         if self._active:
             self._thread_pool.submit(self._trigger, event, *args, **kwargs)
 

ding/framework/middleware/__init__.py

@@ -1,4 +1,7 @@
 from .functional import *
-from .collector import StepCollector, EpisodeCollector
+from .collector import StepCollector, EpisodeCollector, BattleStepCollector
 from .learner import OffPolicyLearner, HERLearner
 from .ckpt_handler import CkptSaver
+from .league_actor import StepLeagueActor
+from .league_coordinator import LeagueCoordinator
+from .league_learner_communicator import LeagueLearnerCommunicator, LearnerModel

ding/framework/middleware/collector.py

@@ -1,13 +1,18 @@
-from typing import TYPE_CHECKING, Callable, List
 from easydict import EasyDict
+from typing import Dict, TYPE_CHECKING, Callable
+import time
+from ditk import logging
 
 from ding.policy import get_random_policy
 from ding.envs import BaseEnvManager
+from ding.utils import log_every_sec
 from ding.framework import task
-from .functional import inferencer, rolloutor, TransitionList
+from ding.framework.middleware.functional import PlayerModelInfo
+from .functional import inferencer, rolloutor, TransitionList, BattleTransitionList, \
+    battle_inferencer, battle_rolloutor
 
 if TYPE_CHECKING:
-    from ding.framework import OnlineRLContext
+    from ding.framework import OnlineRLContext, BattleContext
 
 
 class StepCollector:
@@ -113,4 +118,126 @@ def __call__(self, ctx: "OnlineRLContext") -> None:
                 break
 
 
-# TODO battle collector
+WAIT_MODEL_TIME = float('inf')
+
+
+class BattleStepCollector:
+
+    def __init__(
+        self,
+        cfg: EasyDict,
+        env: BaseEnvManager,
+        unroll_len: int,
+        model_dict: Dict,
+        model_info_dict: Dict,
+        player_policy_collect_dict: Dict,
+        agent_num: int,
+        last_step_fn: Callable = None
+    ):
+        self.cfg = cfg
+        self.end_flag = False
+        # self._reset(env)
+        self.env = env
+        self.env_num = self.env.env_num
+
+        self.total_envstep_count = 0
+        self.unroll_len = unroll_len
+        self.model_dict = model_dict
+        self.model_info_dict = model_info_dict
+        self.player_policy_collect_dict = player_policy_collect_dict
+        self.agent_num = agent_num
+
+        self._battle_inferencer = task.wrap(battle_inferencer(self.cfg, self.env))
+        self._transitions_list = [
+            BattleTransitionList(self.env.env_num, self.unroll_len, last_step_fn) for _ in range(self.agent_num)
+        ]
+        self._battle_rolloutor = task.wrap(
+            battle_rolloutor(self.cfg, self.env, self._transitions_list, self.model_info_dict)
+        )
+
+    def __del__(self) -> None:
+        """
+        Overview:
+            Execute the close command and close the collector. __del__ is automatically called to \
+                destroy the collector instance when the collector finishes its work
+        """
+        if self.end_flag:
+            return
+        self.end_flag = True
+        self.env.close()
+
+    def _update_policies(self, player_id_set) -> None:
+        for player_id in player_id_set:
+            # for this player, if in the beginning of actor's lifetime,
+            # actor didn't recieve any new model, use initial model instead.
+            if self.model_info_dict.get(player_id) is None:
+                self.model_info_dict[player_id] = PlayerModelInfo(
+                    get_new_model_time=time.time(), update_new_model_time=None
+                )
+
+        update_player_id_set = set()
+        for player_id in player_id_set:
+            if 'historical' not in player_id:
+                update_player_id_set.add(player_id)
+        while True:
+            time_now = time.time()
+            time_list = [
+                time_now - self.model_info_dict[player_id].get_new_model_time for player_id in update_player_id_set
+            ]
+            if any(x >= WAIT_MODEL_TIME for x in time_list):
+                for index, player_id in enumerate(update_player_id_set):
+                    if time_list[index] >= WAIT_MODEL_TIME:
+                        log_every_sec(
+                            logging.WARNING, 5,
+                            'In actor {}, model for {} is not updated for {} senconds, and need new model'.format(
+                                task.router.node_id, player_id, time_list[index]
+                            )
+                        )
+                time.sleep(1)
+            else:
+                break
+
+        for player_id in update_player_id_set:
+            if self.model_dict.get(player_id) is None:
+                continue
+            else:
+                learner_model = self.model_dict.get(player_id)
+                policy = self.player_policy_collect_dict.get(player_id)
+                assert policy, "for player{}, policy should have been initialized already"
+                # update policy model
+                policy.load_state_dict(learner_model.state_dict)
+                self.model_info_dict[player_id].update_new_model_time = time.time()
+                self.model_info_dict[player_id].update_train_iter = learner_model.train_iter
+                self.model_dict[player_id] = None
+
+    def __call__(self, ctx: "BattleContext") -> None:
+
+        ctx.total_envstep_count = self.total_envstep_count
+        old = ctx.env_step
+
+        while True:
+            if self.env.closed:
+                self.env.launch()
+                for env_id in range(self.env_num):
+                    for policy in ctx.current_policies:
+                        policy.reset([env_id])
+            self._update_policies(set(ctx.player_id_list))
+            self._battle_inferencer(ctx)
+            self._battle_rolloutor(ctx)
+
+            self.total_envstep_count = ctx.total_envstep_count
+
+            only_finished = True if ctx.env_episode >= ctx.n_episode else False
+            if (self.unroll_len > 0 and ctx.env_step - old >= self.unroll_len) or ctx.env_episode >= ctx.n_episode:
+                for transitions in self._transitions_list:
+                    trajectories = transitions.to_trajectories(only_finished=only_finished)
+                    ctx.trajectories_list.append(trajectories)
+                if ctx.env_episode >= ctx.n_episode:
+                    self.env.close()
+                    ctx.job_finish = True
+                    for transitions in self._transitions_list:
+                        transitions.clear()
+                break
+
+
+# TODO BattleEpisodeCollector

ding/framework/middleware/functional/__init__.py

@@ -1,7 +1,8 @@
 from .trainer import trainer, multistep_trainer
 from .data_processor import offpolicy_data_fetcher, data_pusher, offline_data_fetcher, offline_data_saver, \
     sqil_data_pusher
-from .collector import inferencer, rolloutor, TransitionList
+from .collector import inferencer, rolloutor, TransitionList, BattleTransitionList, \
+    battle_inferencer, battle_rolloutor
 from .evaluator import interaction_evaluator
 from .termination_checker import termination_checker
 from .ctx_helper import final_ctx_saver
@@ -10,3 +11,4 @@
 from .explorer import eps_greedy_handler, eps_greedy_masker
 from .advantage_estimator import gae_estimator
 from .enhancer import reward_estimator, her_data_enhancer, nstep_reward_enhancer
+from .actor_data import ActorData, ActorDataMeta, ActorEnvTrajectories, PlayerModelInfo

ding/framework/middleware/functional/actor_data.py

@@ -0,0 +1,30 @@
+from typing import Any, List
+from dataclasses import dataclass, field
+
+#TODO(zms): simplify fields
+
+
+@dataclass
+class ActorDataMeta:
+    player_total_env_step: int = 0
+    actor_id: int = 0
+    send_wall_time: float = 0.0
+
+
+@dataclass
+class ActorEnvTrajectories:
+    env_id: int = 0
+    trajectories: List = field(default_factory=[])
+
+
+@dataclass
+class ActorData:
+    meta: ActorDataMeta
+    train_data: List[ActorEnvTrajectories] = field(default_factory=[])
+
+
+@dataclass
+class PlayerModelInfo:
+    get_new_model_time: float
+    update_new_model_time: float
+    update_train_iter: int = 0