diff --git a/examples/cim/rl/env_sampler.py b/examples/cim/rl/env_sampler.py index 32c910f36..c7cd241e4 100644 --- a/examples/cim/rl/env_sampler.py +++ b/examples/cim/rl/env_sampler.py @@ -90,11 +90,25 @@ def post_collect(self, info_list: list, ep: int) -> None: for info in info_list: print(f"env summary (episode {ep}): {info['env_metric']}") - # print the average env metric - if len(info_list) > 1: - metric_keys, num_envs = info_list[0]["env_metric"].keys(), len(info_list) - avg_metric = {key: sum(info["env_metric"][key] for info in info_list) / num_envs for key in metric_keys} - print(f"average env summary (episode {ep}): {avg_metric}") + # average env metric + metric_keys, num_envs = info_list[0]["env_metric"].keys(), len(info_list) + avg_metric = {key: sum(info["env_metric"][key] for info in info_list) / num_envs for key in metric_keys} + print(f"average env summary (episode {ep}): {avg_metric}") + + self.metrics.update(avg_metric) + self.metrics = {k: v for k, v in self.metrics.items() if not k.startswith("val/")} def post_evaluate(self, info_list: list, ep: int) -> None: - self.post_collect(info_list, ep) + # print the env metric from each rollout worker + for info in info_list: + print(f"env summary (episode {ep}): {info['env_metric']}") + + # average env metric + metric_keys, num_envs = info_list[0]["env_metric"].keys(), len(info_list) + avg_metric = {key: sum(info["env_metric"][key] for info in info_list) / num_envs for key in metric_keys} + print(f"average env summary (episode {ep}): {avg_metric}") + + self.metrics.update({"val/" + k: v for k, v in avg_metric.items()}) + + def monitor_metrics(self) -> float: + return -self.metrics["val/container_shortage"] diff --git a/examples/cim/rl/rl_component_bundle.py b/examples/cim/rl/rl_component_bundle.py index d290c8f1d..62f6b4fc1 100644 --- a/examples/cim/rl/rl_component_bundle.py +++ b/examples/cim/rl/rl_component_bundle.py @@ -13,7 +13,7 @@ # Environments learn_env = Env(**env_conf) -test_env = learn_env +test_env = Env(**env_conf) # Agent, policy, and trainers num_agents = len(learn_env.agent_idx_list) diff --git a/examples/rl/README.md b/examples/rl/README.md index ca3a3807e..2dc7d2683 100644 --- a/examples/rl/README.md +++ b/examples/rl/README.md @@ -7,7 +7,7 @@ This folder contains scenarios that employ reinforcement learning. MARO's RL too The entrance of a RL workflow is a YAML config file. For readers' convenience, we call this config file `config.yml` in the rest part of this doc. `config.yml` specifies the path of all necessary resources, definitions, and configurations to run the job. MARO provides a comprehensive template of the config file with detailed explanations (`maro/maro/rl/workflows/config/template.yml`). Meanwhile, MARO also provides several simple examples of `config.yml` under the current folder. There are two ways to start the RL job: -- If you only need to have a quick look and try to start an out-of-box workflow, just run `python .\examples\rl\run_rl_example.py PATH_TO_CONFIG_YAML`. For example, `python .\examples\rl\run_rl_example.py .\examples\rl\cim.yml` will run the complete example RL training workflow of CIM scenario. If you only want to run the evaluation workflow, you could start the job with `--evaluate_only`. +- If you only need to have a quick look and try to start an out-of-box workflow, just run `python .\examples\rl\run.py PATH_TO_CONFIG_YAML`. For example, `python .\examples\rl\run.py .\examples\rl\cim.yml` will run the complete example RL training workflow of CIM scenario. If you only want to run the evaluation workflow, you could start the job with `--evaluate_only`. - (**Require install MARO from source**) You could also start the job through MARO CLI. Use the command `maro local run [-c] path/to/your/config` to run in containerized (with `-c`) or non-containerized (without `-c`) environments. Similar, you could add `--evaluate_only` if you only need to run the evaluation workflow. ## Create Your Own Scenarios diff --git a/examples/rl/cim.yml b/examples/rl/cim.yml index 95549a7a8..c383fa9ac 100644 --- a/examples/rl/cim.yml +++ b/examples/rl/cim.yml @@ -10,11 +10,12 @@ job: cim_rl_workflow scenario_path: "examples/cim/rl" -log_path: "log/rl_job/cim.txt" +log_path: "log/cim_rl/" main: num_episodes: 30 # Number of episodes to run. Each episode is one cycle of roll-out and training. num_steps: null eval_schedule: 5 + early_stop_patience: 5 logging: stdout: INFO file: DEBUG @@ -27,7 +28,7 @@ training: load_path: null load_episode: null checkpointing: - path: "checkpoint/rl_job/cim" + path: "log/cim_rl/checkpoints" interval: 5 logging: stdout: INFO diff --git a/examples/rl/cim_distributed.yml b/examples/rl/cim_distributed.yml index 2a52a0846..adbbbc873 100644 --- a/examples/rl/cim_distributed.yml +++ b/examples/rl/cim_distributed.yml @@ -10,7 +10,7 @@ job: cim_rl_workflow scenario_path: "examples/cim/rl" -log_path: "log/rl_job/cim.txt" +log_path: "log/cim_rl/" main: num_episodes: 30 # Number of episodes to run. Each episode is one cycle of roll-out and training. num_steps: null @@ -35,7 +35,7 @@ training: load_path: null load_episode: null checkpointing: - path: "checkpoint/rl_job/cim" + path: "log/cim_rl/checkpoints" interval: 5 proxy: host: "127.0.0.1" diff --git a/examples/rl/vm_scheduling.yml b/examples/rl/vm_scheduling.yml index 7ec2a79e0..16baa67db 100644 --- a/examples/rl/vm_scheduling.yml +++ b/examples/rl/vm_scheduling.yml @@ -10,7 +10,7 @@ job: vm_scheduling_rl_workflow scenario_path: "examples/vm_scheduling/rl" -log_path: "log/rl_job/vm_scheduling.txt" +log_path: "log/vm_rl/" main: num_episodes: 30 # Number of episodes to run. Each episode is one cycle of roll-out and training. num_steps: null @@ -27,7 +27,7 @@ training: load_path: null load_episode: null checkpointing: - path: "checkpoint/rl_job/vm_scheduling" + path: "log/vm_rl/checkpoints" interval: 5 logging: stdout: INFO diff --git a/maro/rl/model/abs_net.py b/maro/rl/model/abs_net.py index a559d1124..0f1430d9c 100644 --- a/maro/rl/model/abs_net.py +++ b/maro/rl/model/abs_net.py @@ -4,7 +4,7 @@ from __future__ import annotations from abc import ABCMeta -from typing import Any, Dict +from typing import Any, Dict, Optional import torch.nn from torch.optim import Optimizer @@ -18,6 +18,8 @@ class AbsNet(torch.nn.Module, metaclass=ABCMeta): def __init__(self) -> None: super(AbsNet, self).__init__() + self._device: Optional[torch.device] = None + @property def optim(self) -> Optimizer: optim = getattr(self, "_optim", None) @@ -119,3 +121,7 @@ def unfreeze_all_parameters(self) -> None: """Unfreeze all parameters.""" for p in self.parameters(): p.requires_grad = True + + def to_device(self, device: torch.device) -> None: + self._device = device + self.to(device) diff --git a/maro/rl/model/algorithm_nets/ac_based.py b/maro/rl/model/algorithm_nets/ac_based.py index 4462cc21d..a4a47a42f 100644 --- a/maro/rl/model/algorithm_nets/ac_based.py +++ b/maro/rl/model/algorithm_nets/ac_based.py @@ -54,3 +54,7 @@ def _get_actions_with_probs_impl(self, states: torch.Tensor, exploring: bool) -> def _get_states_actions_probs_impl(self, states: torch.Tensor, actions: torch.Tensor) -> torch.Tensor: # Not used in Actor-Critic or PPO pass + + def _get_random_actions_impl(self, states: torch.Tensor) -> torch.Tensor: + # Not used in Actor-Critic or PPO + pass diff --git a/maro/rl/model/algorithm_nets/ddpg.py b/maro/rl/model/algorithm_nets/ddpg.py index a4ceb7424..c5e10b009 100644 --- a/maro/rl/model/algorithm_nets/ddpg.py +++ b/maro/rl/model/algorithm_nets/ddpg.py @@ -40,3 +40,7 @@ def _get_states_actions_probs_impl(self, states: torch.Tensor, actions: torch.Te def _get_states_actions_logps_impl(self, states: torch.Tensor, actions: torch.Tensor) -> torch.Tensor: # Not used in DDPG pass + + def _get_random_actions_impl(self, states: torch.Tensor) -> torch.Tensor: + # Not used in DDPG + pass diff --git a/maro/rl/model/policy_net.py b/maro/rl/model/policy_net.py index 9a3b8fd5d..cf1a2df42 100644 --- a/maro/rl/model/policy_net.py +++ b/maro/rl/model/policy_net.py @@ -221,3 +221,17 @@ class ContinuousPolicyNet(PolicyNet, metaclass=ABCMeta): def __init__(self, state_dim: int, action_dim: int) -> None: super(ContinuousPolicyNet, self).__init__(state_dim=state_dim, action_dim=action_dim) + + def get_random_actions(self, states: torch.Tensor) -> torch.Tensor: + actions = self._get_random_actions_impl(states) + + assert self._shape_check( + states=states, + actions=actions, + ), f"Actions shape check failed. Expecting: {(states.shape[0], self.action_dim)}, actual: {actions.shape}." + + return actions + + @abstractmethod + def _get_random_actions_impl(self, states: torch.Tensor) -> torch.Tensor: + raise NotImplementedError diff --git a/maro/rl/policy/abs_policy.py b/maro/rl/policy/abs_policy.py index 14b0bb3a9..d9ea7700c 100644 --- a/maro/rl/policy/abs_policy.py +++ b/maro/rl/policy/abs_policy.py @@ -129,6 +129,8 @@ class RLPolicy(AbsPolicy, metaclass=ABCMeta): state_dim (int): Dimension of states. action_dim (int): Dimension of actions. trainable (bool, default=True): Whether this policy is trainable. + warmup (int, default=0): Number of steps for uniform-random action selection, before running real policy. + Helps exploration. """ def __init__( @@ -138,6 +140,7 @@ def __init__( action_dim: int, is_discrete_action: bool, trainable: bool = True, + warmup: int = 0, ) -> None: super(RLPolicy, self).__init__(name=name, trainable=trainable) self._state_dim = state_dim @@ -145,6 +148,8 @@ def __init__( self._is_exploring = False self._device: Optional[torch.device] = None + self._warmup = warmup + self._call_count = 0 self.is_discrete_action = is_discrete_action @@ -200,7 +205,12 @@ def apply_gradients(self, grad: dict) -> None: raise NotImplementedError def get_actions(self, states: np.ndarray) -> np.ndarray: - actions = self.get_actions_tensor(ndarray_to_tensor(states, device=self._device)) + self._call_count += 1 + + if self._call_count <= self._warmup: + actions = self.get_random_actions_tensor(ndarray_to_tensor(states, device=self._device)) + else: + actions = self.get_actions_tensor(ndarray_to_tensor(states, device=self._device)) return actions.detach().cpu().numpy() def get_actions_tensor(self, states: torch.Tensor) -> torch.Tensor: @@ -217,6 +227,16 @@ def get_actions_tensor(self, states: torch.Tensor) -> torch.Tensor: return actions + def get_random_actions_tensor(self, states: torch.Tensor) -> torch.Tensor: + actions = self._get_random_actions_impl(states) + + assert self._shape_check( + states=states, + actions=actions, + ), f"Actions shape check failed. Expecting: {(states.shape[0], self.action_dim)}, actual: {actions.shape}." + + return actions + def get_actions_with_probs(self, states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: assert self._shape_check( states=states, @@ -273,6 +293,10 @@ def get_states_actions_logps(self, states: torch.Tensor, actions: torch.Tensor) def _get_actions_impl(self, states: torch.Tensor) -> torch.Tensor: raise NotImplementedError + @abstractmethod + def _get_random_actions_impl(self, states: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + @abstractmethod def _get_actions_with_probs_impl(self, states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError diff --git a/maro/rl/policy/continuous_rl_policy.py b/maro/rl/policy/continuous_rl_policy.py index 33ed3e55d..259ba82b1 100644 --- a/maro/rl/policy/continuous_rl_policy.py +++ b/maro/rl/policy/continuous_rl_policy.py @@ -42,6 +42,8 @@ class ContinuousRLPolicy(RLPolicy): the bound for every dimension. If it is a float, it will be broadcast to all dimensions. policy_net (ContinuousPolicyNet): The core net of this policy. trainable (bool, default=True): Whether this policy is trainable. + warmup (int, default=0): Number of steps for uniform-random action selection, before running real policy. + Helps exploration. """ def __init__( @@ -50,6 +52,7 @@ def __init__( action_range: Tuple[Union[float, List[float]], Union[float, List[float]]], policy_net: ContinuousPolicyNet, trainable: bool = True, + warmup: int = 0, ) -> None: assert isinstance(policy_net, ContinuousPolicyNet) @@ -59,6 +62,7 @@ def __init__( action_dim=policy_net.action_dim, trainable=trainable, is_discrete_action=False, + warmup=warmup, ) self._lbounds, self._ubounds = _parse_action_range(self.action_dim, action_range) @@ -83,6 +87,9 @@ def _post_check(self, states: torch.Tensor, actions: torch.Tensor) -> bool: def _get_actions_impl(self, states: torch.Tensor) -> torch.Tensor: return self._policy_net.get_actions(states, self._is_exploring) + def _get_random_actions_impl(self, states: torch.Tensor) -> torch.Tensor: + return self._policy_net.get_random_actions(states) + def _get_actions_with_probs_impl(self, states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: return self._policy_net.get_actions_with_probs(states, self._is_exploring) @@ -117,14 +124,22 @@ def train(self) -> None: self._policy_net.train() def get_state(self) -> dict: - return self._policy_net.get_state() + return { + "net": self._policy_net.get_state(), + "policy": { + "warmup": self._warmup, + "call_count": self._call_count, + }, + } def set_state(self, policy_state: dict) -> None: - self._policy_net.set_state(policy_state) + self._policy_net.set_state(policy_state["net"]) + self._warmup = policy_state["policy"]["warmup"] + self._call_count = policy_state["policy"]["call_count"] def soft_update(self, other_policy: RLPolicy, tau: float) -> None: assert isinstance(other_policy, ContinuousRLPolicy) self._policy_net.soft_update(other_policy.policy_net, tau) def _to_device_impl(self, device: torch.device) -> None: - self._policy_net.to(device) + self._policy_net.to_device(device) diff --git a/maro/rl/policy/discrete_rl_policy.py b/maro/rl/policy/discrete_rl_policy.py index 567e9d054..b2e9d945c 100644 --- a/maro/rl/policy/discrete_rl_policy.py +++ b/maro/rl/policy/discrete_rl_policy.py @@ -23,6 +23,8 @@ class DiscreteRLPolicy(RLPolicy, metaclass=ABCMeta): state_dim (int): Dimension of states. action_num (int): Number of actions. trainable (bool, default=True): Whether this policy is trainable. + warmup (int, default=0): Number of steps for uniform-random action selection, before running real policy. + Helps exploration. """ def __init__( @@ -31,6 +33,7 @@ def __init__( state_dim: int, action_num: int, trainable: bool = True, + warmup: int = 0, ) -> None: assert action_num >= 1 @@ -40,6 +43,7 @@ def __init__( action_dim=1, trainable=trainable, is_discrete_action=True, + warmup=warmup, ) self._action_num = action_num @@ -51,6 +55,12 @@ def action_num(self) -> int: def _post_check(self, states: torch.Tensor, actions: torch.Tensor) -> bool: return all([0 <= action < self.action_num for action in actions.cpu().numpy().flatten()]) + def _get_random_actions_impl(self, states: torch.Tensor) -> torch.Tensor: + return ndarray_to_tensor( + np.random.randint(self.action_num, size=(states.shape[0], 1)), + device=self._device, + ) + class ValueBasedPolicy(DiscreteRLPolicy): """Valued-based policy. @@ -61,7 +71,8 @@ class ValueBasedPolicy(DiscreteRLPolicy): trainable (bool, default=True): Whether this policy is trainable. exploration_strategy (Tuple[Callable, dict], default=(epsilon_greedy, {"epsilon": 0.1})): Exploration strategy. exploration_scheduling_options (List[tuple], default=None): List of exploration scheduler options. - warmup (int, default=50000): Minimum number of experiences to warm up this policy. + warmup (int, default=50000): Number of steps for uniform-random action selection, before running real policy. + Helps exploration. """ def __init__( @@ -80,6 +91,7 @@ def __init__( state_dim=q_net.state_dim, action_num=q_net.action_num, trainable=trainable, + warmup=warmup, ) self._q_net = q_net @@ -91,9 +103,6 @@ def __init__( else [] ) - self._call_cnt = 0 - self._warmup = warmup - self._softmax = torch.nn.Softmax(dim=1) @property @@ -163,19 +172,9 @@ def explore(self) -> None: pass # Overwrite the base method and turn off explore mode. def _get_actions_impl(self, states: torch.Tensor) -> torch.Tensor: - actions, _ = self._get_actions_with_probs_impl(states) - return actions + return self._get_actions_with_probs_impl(states)[0] def _get_actions_with_probs_impl(self, states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - self._call_cnt += 1 - if self._call_cnt <= self._warmup: - actions = ndarray_to_tensor( - np.random.randint(self.action_num, size=(states.shape[0], 1)), - device=self._device, - ) - probs = torch.ones(states.shape[0]).float() * (1.0 / self.action_num) - return actions, probs - q_matrix = self.q_values_for_all_actions_tensor(states) # [B, action_num] q_matrix_softmax = self._softmax(q_matrix) _, actions = q_matrix.max(dim=1) # [B], [B] @@ -222,17 +221,25 @@ def train(self) -> None: self._q_net.train() def get_state(self) -> dict: - return self._q_net.get_state() + return { + "net": self._q_net.get_state(), + "policy": { + "warmup": self._warmup, + "call_count": self._call_count, + }, + } def set_state(self, policy_state: dict) -> None: self._q_net.set_state(policy_state) + self._warmup = policy_state["policy"]["warmup"] + self._call_count = policy_state["policy"]["call_count"] def soft_update(self, other_policy: RLPolicy, tau: float) -> None: assert isinstance(other_policy, ValueBasedPolicy) self._q_net.soft_update(other_policy.q_net, tau) def _to_device_impl(self, device: torch.device) -> None: - self._q_net.to(device) + self._q_net.to_device(device) class DiscretePolicyGradient(DiscreteRLPolicy): @@ -242,6 +249,8 @@ class DiscretePolicyGradient(DiscreteRLPolicy): name (str): Name of the policy. policy_net (DiscretePolicyNet): The core net of this policy. trainable (bool, default=True): Whether this policy is trainable. + warmup (int, default=50000): Number of steps for uniform-random action selection, before running real policy. + Helps exploration. """ def __init__( @@ -249,6 +258,7 @@ def __init__( name: str, policy_net: DiscretePolicyNet, trainable: bool = True, + warmup: int = 0, ) -> None: assert isinstance(policy_net, DiscretePolicyNet) @@ -257,6 +267,7 @@ def __init__( state_dim=policy_net.state_dim, action_num=policy_net.action_num, trainable=trainable, + warmup=warmup, ) self._policy_net = policy_net @@ -302,10 +313,18 @@ def train(self) -> None: self._policy_net.train() def get_state(self) -> dict: - return self._policy_net.get_state() + return { + "net": self._policy_net.get_state(), + "policy": { + "warmup": self._warmup, + "call_count": self._call_count, + }, + } def set_state(self, policy_state: dict) -> None: self._policy_net.set_state(policy_state) + self._warmup = policy_state["policy"]["warmup"] + self._call_count = policy_state["policy"]["call_count"] def soft_update(self, other_policy: RLPolicy, tau: float) -> None: assert isinstance(other_policy, DiscretePolicyGradient) @@ -350,4 +369,4 @@ def _get_state_action_logps_impl(self, states: torch.Tensor, actions: torch.Tens return action_logps.gather(1, actions).squeeze(-1) # [B] def _to_device_impl(self, device: torch.device) -> None: - self._policy_net.to(device) + self._policy_net.to_device(device) diff --git a/maro/rl/rl_component/rl_component_bundle.py b/maro/rl/rl_component/rl_component_bundle.py index f85fe286b..22efbeb8a 100644 --- a/maro/rl/rl_component/rl_component_bundle.py +++ b/maro/rl/rl_component/rl_component_bundle.py @@ -20,7 +20,7 @@ class RLComponentBundle: If None, there will be no explicit device assignment. policy_trainer_mapping (Dict[str, str], default=None): Policy-trainer mapping which identifying which trainer to train each policy. If None, then a policy's trainer's name is the first segment of the policy's name, - seperated by dot. For example, "ppo_1.policy" is trained by "ppo_1". Only policies that provided in + separated by dot. For example, "ppo_1.policy" is trained by "ppo_1". Only policies that provided in policy-trainer mapping are considered as trainable polices. Policies that not provided in policy-trainer mapping will not be trained. """ diff --git a/maro/rl/rollout/batch_env_sampler.py b/maro/rl/rollout/batch_env_sampler.py index a3504a156..b6b0976f0 100644 --- a/maro/rl/rollout/batch_env_sampler.py +++ b/maro/rl/rollout/batch_env_sampler.py @@ -189,8 +189,13 @@ def sample( "info": [res["info"][0] for res in results], } - def eval(self, policy_state: Dict[str, Dict[str, Any]] = None) -> dict: - req = {"type": "eval", "policy_state": policy_state, "index": self._ep} # -1 signals test + def eval(self, policy_state: Dict[str, Dict[str, Any]] = None, num_episodes: int = 1) -> dict: + req = { + "type": "eval", + "policy_state": policy_state, + "index": self._ep, + "num_eval_episodes": num_episodes, + } # -1 signals test results = self._controller.collect(req, self._eval_parallelism) return { "info": [res["info"][0] for res in results], diff --git a/maro/rl/rollout/env_sampler.py b/maro/rl/rollout/env_sampler.py index 81e6ccdad..1e394ac78 100644 --- a/maro/rl/rollout/env_sampler.py +++ b/maro/rl/rollout/env_sampler.py @@ -146,6 +146,7 @@ class ExpElement: terminal_dict: Dict[Any, bool] next_state: Optional[np.ndarray] next_agent_state_dict: Dict[Any, np.ndarray] + truncated: bool @property def agent_names(self) -> list: @@ -171,6 +172,7 @@ def split_contents_by_agent(self) -> Dict[Any, ExpElement]: } if self.next_agent_state_dict is not None and agent_name in self.next_agent_state_dict else {}, + truncated=self.truncated, ) return ret @@ -194,6 +196,7 @@ def split_contents_by_trainer(self, agent2trainer: Dict[Any, str]) -> Dict[str, terminal_dict={}, next_state=self.next_state, next_agent_state_dict=None if self.next_agent_state_dict is None else {}, + truncated=self.truncated, ), ) for agent_name, trainer_name in agent2trainer.items(): @@ -225,6 +228,7 @@ def make_exp_element(self) -> ExpElement: terminal_dict=self.terminal_dict, next_state=self.next_state, next_agent_state_dict=self.next_agent_state_dict, + truncated=self.truncated, ) @@ -240,6 +244,8 @@ class AbsEnvSampler(object, metaclass=ABCMeta): agent_wrapper_cls (Type[AbsAgentWrapper], default=SimpleAgentWrapper): Specific AgentWrapper type. reward_eval_delay (int, default=None): Number of ticks required after a decision event to evaluate the reward for the action taken for that event. If it is None, calculate reward immediately after `step()`. + max_episode_length (int, default=None): Maximum number of steps in one episode during sampling. + When reach this limit, the environment will be truncated and reset. """ def __init__( @@ -251,7 +257,10 @@ def __init__( trainable_policies: List[str] = None, agent_wrapper_cls: Type[AbsAgentWrapper] = SimpleAgentWrapper, reward_eval_delay: int = None, + max_episode_length: int = None, ) -> None: + assert learn_env is not test_env, "Please use different envs for training and testing." + self._learn_env = learn_env self._test_env = test_env @@ -262,11 +271,14 @@ def __init__( self._state: Optional[np.ndarray] = None self._agent_state_dict: Dict[Any, np.ndarray] = {} - self._trans_cache: List[CacheElement] = [] - self._agent_last_index: Dict[Any, int] = {} # Index of last occurrence of agent in self._trans_cache + self._transition_cache: List[CacheElement] = [] + self._agent_last_index: Dict[Any, int] = {} # Index of last occurrence of agent in self._transition_cache self._reward_eval_delay = reward_eval_delay + self._max_episode_length = max_episode_length + self._current_episode_length = 0 self._info: dict = {} + self.metrics: dict = {} assert self._reward_eval_delay is None or self._reward_eval_delay >= 0 @@ -291,11 +303,17 @@ def __init__( [policy_name in self._rl_policy_dict for policy_name in self._trainable_policies], ), "All trainable policies must be RL policies!" + self._total_number_interactions = 0 + @property def env(self) -> Env: assert self._env is not None return self._env + def monitor_metrics(self) -> float: + """Metrics watched by early stopping.""" + return float(self._total_number_interactions) + def _switch_env(self, env: Env) -> None: self._env = env @@ -383,37 +401,37 @@ def _calc_reward(self, cache_element: CacheElement) -> None: def _append_cache_element(self, cache_element: Optional[CacheElement]) -> None: """`cache_element` == None means we are processing the last element in trans_cache""" if cache_element is None: - if len(self._trans_cache) > 0: - self._trans_cache[-1].next_state = self._trans_cache[-1].state - for agent_name, i in self._agent_last_index.items(): - e = self._trans_cache[i] + e = self._transition_cache[i] e.terminal_dict[agent_name] = self._end_of_episode e.next_agent_state_dict[agent_name] = e.agent_state_dict[agent_name] else: - self._trans_cache.append(cache_element) + self._transition_cache.append(cache_element) - if len(self._trans_cache) > 0: - self._trans_cache[-1].next_state = cache_element.state - - cur_index = len(self._trans_cache) - 1 + cur_index = len(self._transition_cache) - 1 for agent_name in cache_element.agent_names: if agent_name in self._agent_last_index: i = self._agent_last_index[agent_name] - self._trans_cache[i].terminal_dict[agent_name] = False - self._trans_cache[i].next_agent_state_dict[agent_name] = cache_element.agent_state_dict[agent_name] + e = self._transition_cache[i] + e.terminal_dict[agent_name] = False + e.next_agent_state_dict[agent_name] = cache_element.agent_state_dict[agent_name] self._agent_last_index[agent_name] = cur_index def _reset(self) -> None: self.env.reset() + self._current_episode_length = 0 self._info.clear() - self._trans_cache.clear() + self._transition_cache.clear() self._agent_last_index.clear() self._step(None) def _select_trainable_agents(self, original_dict: dict) -> dict: return {k: v for k, v in original_dict.items() if k in self._trainable_agents} + @property + def truncated(self) -> bool: + return self._max_episode_length == self._current_episode_length + def sample( self, policy_state: Optional[Dict[str, Dict[str, Any]]] = None, @@ -430,65 +448,88 @@ def sample( Returns: A dict that contains the collected experiences and additional information. """ - # Init the env - self._switch_env(self._learn_env) - if self._end_of_episode: - self._reset() - - # Update policy state if necessary - if policy_state is not None: + steps_to_go = num_steps if num_steps is not None else float("inf") + if policy_state is not None: # Update policy state if necessary self.set_policy_state(policy_state) + self._switch_env(self._learn_env) # Init the env + self._agent_wrapper.explore() # Collect experience - # Collect experience - self._agent_wrapper.explore() - steps_to_go = float("inf") if num_steps is None else num_steps - while not self._end_of_episode and steps_to_go > 0: - # Get agent actions and translate them to env actions - action_dict = self._agent_wrapper.choose_actions(self._agent_state_dict) - env_action_dict = self._translate_to_env_action(action_dict, self._event) - - # Store experiences in the cache - cache_element = CacheElement( - tick=self.env.tick, - event=self._event, - state=self._state, - agent_state_dict=self._select_trainable_agents(self._agent_state_dict), - action_dict=self._select_trainable_agents(action_dict), - env_action_dict=self._select_trainable_agents(env_action_dict), - # The following will be generated later - reward_dict={}, - terminal_dict={}, - next_state=None, - next_agent_state_dict={}, - ) + if self._end_of_episode: + self._reset() - # Update env and get new states (global & agent) - self._step(list(env_action_dict.values())) - - if self._reward_eval_delay is None: - self._calc_reward(cache_element) - self._post_step(cache_element) - self._append_cache_element(cache_element) - steps_to_go -= 1 - self._append_cache_element(None) - - tick_bound = self.env.tick - (0 if self._reward_eval_delay is None else self._reward_eval_delay) - experiences: List[ExpElement] = [] - while len(self._trans_cache) > 0 and self._trans_cache[0].tick <= tick_bound: - cache_element = self._trans_cache.pop(0) - # !: Here the reward calculation method requires the given tick is enough and must be used then. - if self._reward_eval_delay is not None: - self._calc_reward(cache_element) - self._post_step(cache_element) - experiences.append(cache_element.make_exp_element()) - - self._agent_last_index = { - k: v - len(experiences) for k, v in self._agent_last_index.items() if v >= len(experiences) - } + # If num_steps is None, run until the end of episode or the episode is truncated + # If num_steps is not None, run until we collect required number of steps + total_experiences = [] + + while not any( + [ + num_steps is None and (self._end_of_episode or self.truncated), + num_steps is not None and steps_to_go == 0, + ], + ): + if self._end_of_episode or self.truncated: + self._reset() + + while not any( + [ + self._end_of_episode, + self.truncated, + steps_to_go == 0, + ], + ): + # Get agent actions and translate them to env actions + action_dict = self._agent_wrapper.choose_actions(self._agent_state_dict) + env_action_dict = self._translate_to_env_action(action_dict, self._event) + + self._total_number_interactions += 1 + self._current_episode_length += 1 + steps_to_go -= 1 + + # Store experiences in the cache + cache_element = CacheElement( + tick=self.env.tick, + event=self._event, + state=self._state, + agent_state_dict=self._select_trainable_agents(self._agent_state_dict), + action_dict=self._select_trainable_agents(action_dict), + env_action_dict=self._select_trainable_agents(env_action_dict), + # The following will be generated/updated later + reward_dict={}, + terminal_dict={}, + next_state=None, + next_agent_state_dict={}, + truncated=self.truncated, + ) + + # Update env and get new states (global & agent) + self._step(list(env_action_dict.values())) + cache_element.next_state = self._state + + if self._reward_eval_delay is None: + self._calc_reward(cache_element) + self._post_step(cache_element) + self._append_cache_element(cache_element) + + self._append_cache_element(None) + + tick_bound = self.env.tick - (0 if self._reward_eval_delay is None else self._reward_eval_delay) + experiences: List[ExpElement] = [] + while len(self._transition_cache) > 0 and self._transition_cache[0].tick <= tick_bound: + cache_element = self._transition_cache.pop(0) + # !: Here the reward calculation method requires the given tick is enough and must be used then. + if self._reward_eval_delay is not None: + self._calc_reward(cache_element) + self._post_step(cache_element) + experiences.append(cache_element.make_exp_element()) + + self._agent_last_index = { + k: v - len(experiences) for k, v in self._agent_last_index.items() if v >= len(experiences) + } + + total_experiences += experiences return { - "end_of_episode": self._end_of_episode, - "experiences": [experiences], + "experiences": [total_experiences], "info": [deepcopy(self._info)], # TODO: may have overhead issues. Leave to future work. } @@ -514,50 +555,57 @@ def load_policy_state(self, path: str) -> List[str]: return loaded - def eval(self, policy_state: Dict[str, Dict[str, Any]] = None) -> dict: + def eval(self, policy_state: Dict[str, Dict[str, Any]] = None, num_episodes: int = 1) -> dict: self._switch_env(self._test_env) - self._reset() - if policy_state is not None: - self.set_policy_state(policy_state) + info_list = [] - self._agent_wrapper.exploit() - while not self._end_of_episode: - action_dict = self._agent_wrapper.choose_actions(self._agent_state_dict) - env_action_dict = self._translate_to_env_action(action_dict, self._event) - - # Store experiences in the cache - cache_element = CacheElement( - tick=self.env.tick, - event=self._event, - state=self._state, - agent_state_dict=self._select_trainable_agents(self._agent_state_dict), - action_dict=self._select_trainable_agents(action_dict), - env_action_dict=self._select_trainable_agents(env_action_dict), - # The following will be generated later - reward_dict={}, - terminal_dict={}, - next_state=None, - next_agent_state_dict={}, - ) - - # Update env and get new states (global & agent) - self._step(list(env_action_dict.values())) - - if self._reward_eval_delay is None: # TODO: necessary to calculate reward in eval()? - self._calc_reward(cache_element) - self._post_eval_step(cache_element) - - self._append_cache_element(cache_element) - self._append_cache_element(None) - - tick_bound = self.env.tick - (0 if self._reward_eval_delay is None else self._reward_eval_delay) - while len(self._trans_cache) > 0 and self._trans_cache[0].tick <= tick_bound: - cache_element = self._trans_cache.pop(0) - if self._reward_eval_delay is not None: - self._calc_reward(cache_element) - self._post_eval_step(cache_element) - - return {"info": [self._info]} + for _ in range(num_episodes): + self._reset() + if policy_state is not None: + self.set_policy_state(policy_state) + + self._agent_wrapper.exploit() + while not self._end_of_episode: + action_dict = self._agent_wrapper.choose_actions(self._agent_state_dict) + env_action_dict = self._translate_to_env_action(action_dict, self._event) + + # Store experiences in the cache + cache_element = CacheElement( + tick=self.env.tick, + event=self._event, + state=self._state, + agent_state_dict=self._select_trainable_agents(self._agent_state_dict), + action_dict=self._select_trainable_agents(action_dict), + env_action_dict=self._select_trainable_agents(env_action_dict), + # The following will be generated later + reward_dict={}, + terminal_dict={}, + next_state=None, + next_agent_state_dict={}, + truncated=False, # No truncation in evaluation + ) + + # Update env and get new states (global & agent) + self._step(list(env_action_dict.values())) + cache_element.next_state = self._state + + if self._reward_eval_delay is None: # TODO: necessary to calculate reward in eval()? + self._calc_reward(cache_element) + self._post_eval_step(cache_element) + + self._append_cache_element(cache_element) + self._append_cache_element(None) + + tick_bound = self.env.tick - (0 if self._reward_eval_delay is None else self._reward_eval_delay) + while len(self._transition_cache) > 0 and self._transition_cache[0].tick <= tick_bound: + cache_element = self._transition_cache.pop(0) + if self._reward_eval_delay is not None: + self._calc_reward(cache_element) + self._post_eval_step(cache_element) + + info_list.append(self._info) + + return {"info": info_list} @abstractmethod def _post_step(self, cache_element: CacheElement) -> None: diff --git a/maro/rl/rollout/worker.py b/maro/rl/rollout/worker.py index b8301ee38..1532a6489 100644 --- a/maro/rl/rollout/worker.py +++ b/maro/rl/rollout/worker.py @@ -59,7 +59,7 @@ def _compute(self, msg: list) -> None: result = ( self._env_sampler.sample(policy_state=req["policy_state"], num_steps=req["num_steps"]) if req["type"] == "sample" - else self._env_sampler.eval(policy_state=req["policy_state"]) + else self._env_sampler.eval(policy_state=req["policy_state"], num_episodes=req["num_eval_episodes"]) ) self._stream.send(pyobj_to_bytes({"result": result, "index": req["index"]})) else: diff --git a/maro/rl/training/algorithms/base/ac_ppo_base.py b/maro/rl/training/algorithms/base/ac_ppo_base.py index 3227437be..aeead3574 100644 --- a/maro/rl/training/algorithms/base/ac_ppo_base.py +++ b/maro/rl/training/algorithms/base/ac_ppo_base.py @@ -202,6 +202,9 @@ def preprocess_batch(self, batch: TransitionBatch) -> TransitionBatch: # Preprocess advantages states = ndarray_to_tensor(batch.states, device=self._device) # s actions = ndarray_to_tensor(batch.actions, device=self._device) # a + terminals = ndarray_to_tensor(batch.terminals, device=self._device) + truncated = ndarray_to_tensor(batch.truncated, device=self._device) + next_states = ndarray_to_tensor(batch.next_states, device=self._device) if self._is_discrete_action: actions = actions.long() @@ -209,11 +212,34 @@ def preprocess_batch(self, batch: TransitionBatch) -> TransitionBatch: self._v_critic_net.eval() self._policy.eval() values = self._v_critic_net.v_values(states).detach().cpu().numpy() - values = np.concatenate([values, np.zeros(1)]) - rewards = np.concatenate([batch.rewards, np.zeros(1)]) - deltas = rewards[:-1] + self._reward_discount * values[1:] - values[:-1] # r + gamma * v(s') - v(s) - batch.returns = discount_cumsum(rewards, self._reward_discount)[:-1] - batch.advantages = discount_cumsum(deltas, self._reward_discount * self._lam) + + batch.returns = np.zeros(batch.size, dtype=np.float32) + batch.advantages = np.zeros(batch.size, dtype=np.float32) + i = 0 + while i < batch.size: + j = i + while j < batch.size - 1 and not (terminals[j] or truncated[j]): + j += 1 + last_val = ( + 0.0 + if terminals[j] + else self._v_critic_net.v_values( + next_states[j].unsqueeze(dim=0), + ) + .detach() + .cpu() + .numpy() + .item() + ) + + cur_values = np.append(values[i : j + 1], last_val) + cur_rewards = np.append(batch.rewards[i : j + 1], last_val) + # delta = r + gamma * v(s') - v(s) + cur_deltas = cur_rewards[:-1] + self._reward_discount * cur_values[1:] - cur_values[:-1] + batch.returns[i : j + 1] = discount_cumsum(cur_rewards, self._reward_discount)[:-1] + batch.advantages[i : j + 1] = discount_cumsum(cur_deltas, self._reward_discount * self._lam) + + i = j + 1 if self._clip_ratio is not None: batch.old_logps = self._policy.get_states_actions_logps(states, actions).detach().cpu().numpy() @@ -291,21 +317,23 @@ def train_step(self) -> None: assert isinstance(self._ops, ACBasedOps) batch = self._get_batch() - for _ in range(self._params.grad_iters): - self._ops.update_critic(batch) for _ in range(self._params.grad_iters): early_stop = self._ops.update_actor(batch) if early_stop: break + for _ in range(self._params.grad_iters): + self._ops.update_critic(batch) + async def train_step_as_task(self) -> None: assert isinstance(self._ops, RemoteOps) batch = self._get_batch() - for _ in range(self._params.grad_iters): - self._ops.update_critic_with_grad(await self._ops.get_critic_grad(batch)) for _ in range(self._params.grad_iters): if self._ops.update_actor_with_grad(await self._ops.get_actor_grad(batch)): # early stop break + + for _ in range(self._params.grad_iters): + self._ops.update_critic_with_grad(await self._ops.get_critic_grad(batch)) diff --git a/maro/rl/training/algorithms/ddpg.py b/maro/rl/training/algorithms/ddpg.py index 79bd5b336..53b070300 100644 --- a/maro/rl/training/algorithms/ddpg.py +++ b/maro/rl/training/algorithms/ddpg.py @@ -27,7 +27,7 @@ class DDPGParams(BaseTrainerParams): random_overwrite (bool, default=False): This specifies overwrite behavior when the replay memory capacity is reached. If True, overwrite positions will be selected randomly. Otherwise, overwrites will occur sequentially with wrap-around. - min_num_to_trigger_training (int, default=0): Minimum number required to start training. + n_start_train (int, default=0): Minimum number required to start training. """ get_q_critic_net_func: Callable[[], QNet] @@ -36,7 +36,7 @@ class DDPGParams(BaseTrainerParams): q_value_loss_cls: Optional[Callable] = None soft_update_coef: float = 1.0 random_overwrite: bool = False - min_num_to_trigger_training: int = 0 + n_start_train: int = 0 class DDPGOps(AbsTrainOps): @@ -93,9 +93,9 @@ def _get_critic_loss(self, batch: TransitionBatch) -> torch.Tensor: states=next_states, # s' actions=self._target_policy.get_actions_tensor(next_states), # miu_targ(s') ) # Q_targ(s', miu_targ(s')) + # y(r, s', d) = r + gamma * (1 - d) * Q_targ(s', miu_targ(s')) + target_q_values = (rewards + self._reward_discount * (1.0 - terminals.float()) * next_q_values).detach() - # y(r, s', d) = r + gamma * (1 - d) * Q_targ(s', miu_targ(s')) - target_q_values = (rewards + self._reward_discount * (1 - terminals.long()) * next_q_values).detach() q_values = self._q_critic_net.q_values(states=states, actions=actions) # Q(s, a) return self._q_value_loss_func(q_values, target_q_values) # MSE(Q(s, a), y(r, s', d)) @@ -263,10 +263,10 @@ def _get_batch(self, batch_size: int = None) -> TransitionBatch: def train_step(self) -> None: assert isinstance(self._ops, DDPGOps) - if self._replay_memory.n_sample < self._params.min_num_to_trigger_training: + if self._replay_memory.n_sample < self._params.n_start_train: print( f"Skip this training step due to lack of experiences " - f"(current = {self._replay_memory.n_sample}, minimum = {self._params.min_num_to_trigger_training})", + f"(current = {self._replay_memory.n_sample}, minimum = {self._params.n_start_train})", ) return @@ -280,10 +280,10 @@ def train_step(self) -> None: async def train_step_as_task(self) -> None: assert isinstance(self._ops, RemoteOps) - if self._replay_memory.n_sample < self._params.min_num_to_trigger_training: + if self._replay_memory.n_sample < self._params.n_start_train: print( f"Skip this training step due to lack of experiences " - f"(current = {self._replay_memory.n_sample}, minimum = {self._params.min_num_to_trigger_training})", + f"(current = {self._replay_memory.n_sample}, minimum = {self._params.n_start_train})", ) return diff --git a/maro/rl/training/algorithms/maddpg.py b/maro/rl/training/algorithms/maddpg.py index edc63f39a..1e5d1d766 100644 --- a/maro/rl/training/algorithms/maddpg.py +++ b/maro/rl/training/algorithms/maddpg.py @@ -378,6 +378,7 @@ def record_multiple(self, env_idx: int, exp_elements: List[ExpElement]) -> None: agent_states=agent_states, next_agent_states=next_agent_states, terminals=np.array(terminal_flags), + truncated=np.array([exp_element.truncated for exp_element in exp_elements]), ) self._replay_memory.put(transition_batch) diff --git a/maro/rl/training/algorithms/sac.py b/maro/rl/training/algorithms/sac.py index 338addf57..d7332da7e 100644 --- a/maro/rl/training/algorithms/sac.py +++ b/maro/rl/training/algorithms/sac.py @@ -22,7 +22,7 @@ class SoftActorCriticParams(BaseTrainerParams): num_epochs: int = 1 n_start_train: int = 0 q_value_loss_cls: Optional[Callable] = None - soft_update_coef: float = 1.0 + soft_update_coef: float = 0.05 class SoftActorCriticOps(AbsTrainOps): @@ -58,6 +58,7 @@ def __init__( def _get_critic_loss(self, batch: TransitionBatch) -> Tuple[torch.Tensor, torch.Tensor]: self._q_net1.train() + self._q_net2.train() states = ndarray_to_tensor(batch.states, device=self._device) # s next_states = ndarray_to_tensor(batch.next_states, device=self._device) # s' actions = ndarray_to_tensor(batch.actions, device=self._device) # a @@ -67,11 +68,13 @@ def _get_critic_loss(self, batch: TransitionBatch) -> Tuple[torch.Tensor, torch. assert isinstance(self._policy, ContinuousRLPolicy) with torch.no_grad(): - next_actions, next_logps = self._policy.get_actions_with_logps(states) - q1 = self._target_q_net1.q_values(next_states, next_actions) - q2 = self._target_q_net2.q_values(next_states, next_actions) - q = torch.min(q1, q2) - y = rewards + self._reward_discount * (1.0 - terminals.float()) * (q - self._entropy_coef * next_logps) + next_actions, next_logps = self._policy.get_actions_with_logps(next_states) + target_q1 = self._target_q_net1.q_values(next_states, next_actions) + target_q2 = self._target_q_net2.q_values(next_states, next_actions) + target_q = torch.min(target_q1, target_q2) + y = rewards + self._reward_discount * (1.0 - terminals.float()) * ( + target_q - self._entropy_coef * next_logps + ) q1 = self._q_net1.q_values(states, actions) q2 = self._q_net2.q_values(states, actions) @@ -100,6 +103,9 @@ def update_critic(self, batch: TransitionBatch) -> None: self._q_net2.step(loss_q2) def _get_actor_loss(self, batch: TransitionBatch) -> torch.Tensor: + self._q_net1.freeze() + self._q_net2.freeze() + self._policy.train() states = ndarray_to_tensor(batch.states, device=self._device) # s actions, logps = self._policy.get_actions_with_logps(states) @@ -108,6 +114,10 @@ def _get_actor_loss(self, batch: TransitionBatch) -> torch.Tensor: q = torch.min(q1, q2) loss = (self._entropy_coef * logps - q).mean() + + self._q_net1.unfreeze() + self._q_net2.unfreeze() + return loss @remote @@ -142,6 +152,9 @@ def soft_update_target(self) -> None: def to_device(self, device: str = None) -> None: self._device = get_torch_device(device=device) + + self._policy.to_device(self._device) + self._q_net1.to(self._device) self._q_net2.to(self._device) self._target_q_net1.to(self._device) diff --git a/maro/rl/training/replay_memory.py b/maro/rl/training/replay_memory.py index 3e4f573e0..da1e7d692 100644 --- a/maro/rl/training/replay_memory.py +++ b/maro/rl/training/replay_memory.py @@ -35,29 +35,18 @@ def get_put_indexes(self, batch_size: int) -> np.ndarray: raise NotImplementedError @abstractmethod - def get_sample_indexes(self, batch_size: int = None, forbid_last: bool = False) -> np.ndarray: + def get_sample_indexes(self, batch_size: int = None) -> np.ndarray: """Generate a list of indexes that can be used to retrieve items from the replay memory. Args: batch_size (int, default=None): The required batch size. If it is None, all indexes where an experience item is present are returned. - forbid_last (bool, default=False): Whether the latest element is allowed to be sampled. - If this is true, the last index will always be excluded from the result. Returns: indexes (np.ndarray): The list of indexes. """ raise NotImplementedError - @abstractmethod - def get_last_index(self) -> int: - """Get the index of the latest element in the memory. - - Returns: - index (int): The index of the latest element in the memory. - """ - raise NotImplementedError - class RandomIndexScheduler(AbsIndexScheduler): """Index scheduler that returns random indexes when sampling. @@ -93,14 +82,11 @@ def get_put_indexes(self, batch_size: int) -> np.ndarray: self._size = min(self._size + batch_size, self._capacity) return indexes - def get_sample_indexes(self, batch_size: int = None, forbid_last: bool = False) -> np.ndarray: + def get_sample_indexes(self, batch_size: int = None) -> np.ndarray: assert batch_size is not None and batch_size > 0, f"Invalid batch size: {batch_size}" assert self._size > 0, "Cannot sample from an empty memory." return np.random.choice(self._size, size=batch_size, replace=True) - def get_last_index(self) -> int: - raise NotImplementedError - class FIFOIndexScheduler(AbsIndexScheduler): """First-in-first-out index scheduler. @@ -135,19 +121,15 @@ def get_put_indexes(self, batch_size: int) -> np.ndarray: self._head = (self._head + overwrite) % self._capacity return self.get_put_indexes(batch_size) - def get_sample_indexes(self, batch_size: int = None, forbid_last: bool = False) -> np.ndarray: - tmp = self._tail if not forbid_last else (self._tail - 1) % self._capacity + def get_sample_indexes(self, batch_size: int = None) -> np.ndarray: indexes = ( - np.arange(self._head, tmp) - if tmp > self._head - else np.concatenate([np.arange(self._head, self._capacity), np.arange(tmp)]) + np.arange(self._head, self._tail) + if self._tail > self._head + else np.concatenate([np.arange(self._head, self._capacity), np.arange(self._tail)]) ) - self._head = tmp + self._head = self._tail return indexes - def get_last_index(self) -> int: - return (self._tail - 1) % self._capacity - class AbsReplayMemory(object, metaclass=ABCMeta): """Abstract replay memory class with basic interfaces. @@ -176,9 +158,9 @@ def _get_put_indexes(self, batch_size: int) -> np.ndarray: """Please refer to the doc string in AbsIndexScheduler.""" return self._idx_scheduler.get_put_indexes(batch_size) - def _get_sample_indexes(self, batch_size: int = None, forbid_last: bool = False) -> np.ndarray: + def _get_sample_indexes(self, batch_size: int = None) -> np.ndarray: """Please refer to the doc string in AbsIndexScheduler.""" - return self._idx_scheduler.get_sample_indexes(batch_size, forbid_last) + return self._idx_scheduler.get_sample_indexes(batch_size) class ReplayMemory(AbsReplayMemory, metaclass=ABCMeta): @@ -205,6 +187,7 @@ def __init__( self._actions = np.zeros((self._capacity, self._action_dim), dtype=np.float32) self._rewards = np.zeros(self._capacity, dtype=np.float32) self._terminals = np.zeros(self._capacity, dtype=bool) + self._truncated = np.zeros(self._capacity, dtype=bool) self._next_states = np.zeros((self._capacity, self._state_dim), dtype=np.float32) self._returns = np.zeros(self._capacity, dtype=np.float32) self._advantages = np.zeros(self._capacity, dtype=np.float32) @@ -233,6 +216,7 @@ def put(self, transition_batch: TransitionBatch) -> None: assert match_shape(transition_batch.actions, (batch_size, self._action_dim)) assert match_shape(transition_batch.rewards, (batch_size,)) assert match_shape(transition_batch.terminals, (batch_size,)) + assert match_shape(transition_batch.truncated, (batch_size,)) assert match_shape(transition_batch.next_states, (batch_size, self._state_dim)) if transition_batch.returns is not None: match_shape(transition_batch.returns, (batch_size,)) @@ -255,6 +239,7 @@ def _put_by_indexes(self, indexes: np.ndarray, transition_batch: TransitionBatch self._actions[indexes] = transition_batch.actions self._rewards[indexes] = transition_batch.rewards self._terminals[indexes] = transition_batch.terminals + self._truncated[indexes] = transition_batch.truncated self._next_states[indexes] = transition_batch.next_states if transition_batch.returns is not None: self._returns[indexes] = transition_batch.returns @@ -273,7 +258,7 @@ def sample(self, batch_size: int = None) -> TransitionBatch: Returns: batch (TransitionBatch): The sampled batch. """ - indexes = self._get_sample_indexes(batch_size, self._get_forbid_last()) + indexes = self._get_sample_indexes(batch_size) return self.sample_by_indexes(indexes) def sample_by_indexes(self, indexes: np.ndarray) -> TransitionBatch: @@ -292,16 +277,13 @@ def sample_by_indexes(self, indexes: np.ndarray) -> TransitionBatch: actions=self._actions[indexes], rewards=self._rewards[indexes], terminals=self._terminals[indexes], + truncated=self._truncated[indexes], next_states=self._next_states[indexes], returns=self._returns[indexes], advantages=self._advantages[indexes], old_logps=self._old_logps[indexes], ) - @abstractmethod - def _get_forbid_last(self) -> bool: - raise NotImplementedError - class RandomReplayMemory(ReplayMemory): def __init__( @@ -318,15 +300,11 @@ def __init__( RandomIndexScheduler(capacity, random_overwrite), ) self._random_overwrite = random_overwrite - self._scheduler = RandomIndexScheduler(capacity, random_overwrite) @property def random_overwrite(self) -> bool: return self._random_overwrite - def _get_forbid_last(self) -> bool: - return False - class FIFOReplayMemory(ReplayMemory): def __init__( @@ -342,9 +320,6 @@ def __init__( FIFOIndexScheduler(capacity), ) - def _get_forbid_last(self) -> bool: - return not self._terminals[self._idx_scheduler.get_last_index()] - class MultiReplayMemory(AbsReplayMemory, metaclass=ABCMeta): """In-memory experience storage facility for a multi trainer. @@ -374,6 +349,7 @@ def __init__( self._rewards = [np.zeros(self._capacity, dtype=np.float32) for _ in range(self.agent_num)] self._next_states = np.zeros((self._capacity, self._state_dim), dtype=np.float32) self._terminals = np.zeros(self._capacity, dtype=bool) + self._truncated = np.zeros(self._capacity, dtype=bool) assert len(agent_states_dims) == self.agent_num self._agent_states_dims = agent_states_dims @@ -408,6 +384,7 @@ def put(self, transition_batch: MultiTransitionBatch) -> None: assert match_shape(transition_batch.rewards[i], (batch_size,)) assert match_shape(transition_batch.terminals, (batch_size,)) + assert match_shape(transition_batch.truncated, (batch_size,)) assert match_shape(transition_batch.next_states, (batch_size, self._state_dim)) assert len(transition_batch.agent_states) == self.agent_num @@ -430,6 +407,7 @@ def _put_by_indexes(self, indexes: np.ndarray, transition_batch: MultiTransition self._actions[i][indexes] = transition_batch.actions[i] self._rewards[i][indexes] = transition_batch.rewards[i] self._terminals[indexes] = transition_batch.terminals + self._truncated[indexes] = transition_batch.truncated self._next_states[indexes] = transition_batch.next_states for i in range(self.agent_num): @@ -446,7 +424,7 @@ def sample(self, batch_size: int = None) -> MultiTransitionBatch: Returns: batch (MultiTransitionBatch): The sampled batch. """ - indexes = self._get_sample_indexes(batch_size, self._get_forbid_last()) + indexes = self._get_sample_indexes(batch_size) return self.sample_by_indexes(indexes) def sample_by_indexes(self, indexes: np.ndarray) -> MultiTransitionBatch: @@ -465,15 +443,12 @@ def sample_by_indexes(self, indexes: np.ndarray) -> MultiTransitionBatch: actions=[action[indexes] for action in self._actions], rewards=[reward[indexes] for reward in self._rewards], terminals=self._terminals[indexes], + truncated=self._truncated[indexes], next_states=self._next_states[indexes], agent_states=[state[indexes] for state in self._agent_states], next_agent_states=[state[indexes] for state in self._next_agent_states], ) - @abstractmethod - def _get_forbid_last(self) -> bool: - raise NotImplementedError - class RandomMultiReplayMemory(MultiReplayMemory): def __init__( @@ -492,15 +467,11 @@ def __init__( agent_states_dims, ) self._random_overwrite = random_overwrite - self._scheduler = RandomIndexScheduler(capacity, random_overwrite) @property def random_overwrite(self) -> bool: return self._random_overwrite - def _get_forbid_last(self) -> bool: - return False - class FIFOMultiReplayMemory(MultiReplayMemory): def __init__( @@ -517,6 +488,3 @@ def __init__( FIFOIndexScheduler(capacity), agent_states_dims, ) - - def _get_forbid_last(self) -> bool: - return not self._terminals[self._idx_scheduler.get_last_index()] diff --git a/maro/rl/training/trainer.py b/maro/rl/training/trainer.py index 8bced5674..774954f6c 100644 --- a/maro/rl/training/trainer.py +++ b/maro/rl/training/trainer.py @@ -254,6 +254,7 @@ def record_multiple(self, env_idx: int, exp_elements: List[ExpElement]) -> None: exp_element.action_dict[agent_name], exp_element.reward_dict[agent_name], exp_element.terminal_dict[agent_name], + exp_element.truncated, exp_element.next_agent_state_dict.get(agent_name, exp_element.agent_state_dict[agent_name]), ), ) @@ -264,7 +265,8 @@ def record_multiple(self, env_idx: int, exp_elements: List[ExpElement]) -> None: actions=np.vstack([exp[1] for exp in exps]), rewards=np.array([exp[2] for exp in exps]), terminals=np.array([exp[3] for exp in exps]), - next_states=np.vstack([exp[4] for exp in exps]), + truncated=np.array([exp[4] for exp in exps]), + next_states=np.vstack([exp[5] for exp in exps]), ) transition_batch = self._preprocess_batch(transition_batch) self.replay_memory.put(transition_batch) diff --git a/maro/rl/utils/transition_batch.py b/maro/rl/utils/transition_batch.py index f9ada5473..53bbd6233 100644 --- a/maro/rl/utils/transition_batch.py +++ b/maro/rl/utils/transition_batch.py @@ -19,6 +19,7 @@ class TransitionBatch: rewards: np.ndarray # 1D next_states: np.ndarray # 2D terminals: np.ndarray # 1D + truncated: np.ndarray # 1D returns: np.ndarray = None # 1D advantages: np.ndarray = None # 1D old_logps: np.ndarray = None # 1D @@ -34,6 +35,7 @@ def __post_init__(self) -> None: assert len(self.rewards.shape) == 1 and self.rewards.shape[0] == self.states.shape[0] assert self.next_states.shape == self.states.shape assert len(self.terminals.shape) == 1 and self.terminals.shape[0] == self.states.shape[0] + assert len(self.truncated.shape) == 1 and self.truncated.shape[0] == self.states.shape[0] def make_kth_sub_batch(self, i: int, k: int) -> TransitionBatch: return TransitionBatch( @@ -42,6 +44,7 @@ def make_kth_sub_batch(self, i: int, k: int) -> TransitionBatch: rewards=self.rewards[i::k], next_states=self.next_states[i::k], terminals=self.terminals[i::k], + truncated=self.truncated[i::k], returns=self.returns[i::k] if self.returns is not None else None, advantages=self.advantages[i::k] if self.advantages is not None else None, old_logps=self.old_logps[i::k] if self.old_logps is not None else None, @@ -60,7 +63,7 @@ class MultiTransitionBatch: agent_states: List[np.ndarray] # List of 2D next_agent_states: List[np.ndarray] # List of 2D terminals: np.ndarray # 1D - + truncated: np.ndarray # 1D returns: Optional[List[np.ndarray]] = None # List of 1D advantages: Optional[List[np.ndarray]] = None # List of 1D @@ -81,6 +84,7 @@ def __post_init__(self) -> None: assert self.agent_states[i].shape[0] == self.states.shape[0] assert len(self.terminals.shape) == 1 and self.terminals.shape[0] == self.states.shape[0] + assert len(self.truncated.shape) == 1 and self.truncated.shape[0] == self.states.shape[0] assert self.next_states.shape == self.states.shape assert len(self.next_agent_states) == len(self.agent_states) @@ -98,6 +102,7 @@ def make_kth_sub_batch(self, i: int, k: int) -> MultiTransitionBatch: agent_states = [state[i::k] for state in self.agent_states] next_agent_states = [state[i::k] for state in self.next_agent_states] terminals = self.terminals[i::k] + truncated = self.truncated[i::k] returns = None if self.returns is None else [r[i::k] for r in self.returns] advantages = None if self.advantages is None else [advantage[i::k] for advantage in self.advantages] return MultiTransitionBatch( @@ -108,6 +113,7 @@ def make_kth_sub_batch(self, i: int, k: int) -> MultiTransitionBatch: agent_states, next_agent_states, terminals, + truncated, returns, advantages, ) @@ -123,6 +129,7 @@ def merge_transition_batches(batch_list: List[TransitionBatch]) -> TransitionBat rewards=np.concatenate([batch.rewards for batch in batch_list], axis=0), next_states=np.concatenate([batch.next_states for batch in batch_list], axis=0), terminals=np.concatenate([batch.terminals for batch in batch_list]), + truncated=np.concatenate([batch.truncated for batch in batch_list]), returns=np.concatenate([batch.returns for batch in batch_list]), advantages=np.concatenate([batch.advantages for batch in batch_list]), old_logps=None diff --git a/maro/rl/workflows/callback.py b/maro/rl/workflows/callback.py new file mode 100644 index 000000000..1c5a2c2f7 --- /dev/null +++ b/maro/rl/workflows/callback.py @@ -0,0 +1,182 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from __future__ import annotations + +import copy +import os +import typing +from typing import Dict, List, Optional, Union + +import pandas as pd + +from maro.rl.rollout import AbsEnvSampler, BatchEnvSampler +from maro.rl.training import TrainingManager +from maro.utils import LoggerV2 + +if typing.TYPE_CHECKING: + from maro.rl.workflows.main import TrainingWorkflow + +EnvSampler = Union[AbsEnvSampler, BatchEnvSampler] + + +class Callback(object): + def __init__(self) -> None: + self.workflow: Optional[TrainingWorkflow] = None + self.env_sampler: Optional[EnvSampler] = None + self.training_manager: Optional[TrainingManager] = None + self.logger: Optional[LoggerV2] = None + + def on_episode_start(self, ep: int) -> None: + pass + + def on_episode_end(self, ep: int) -> None: + pass + + def on_training_start(self, ep: int) -> None: + pass + + def on_training_end(self, ep: int) -> None: + pass + + def on_validation_start(self, ep: int) -> None: + pass + + def on_validation_end(self, ep: int) -> None: + pass + + def on_test_start(self, ep: int) -> None: + pass + + def on_test_end(self, ep: int) -> None: + pass + + +class EarlyStopping(Callback): + def __init__(self, patience: int) -> None: + super(EarlyStopping, self).__init__() + + self._patience = patience + self._best_ep: int = -1 + self._best: float = float("-inf") + + def on_validation_end(self, ep: int) -> None: + cur = self.env_sampler.monitor_metrics() + if cur > self._best: + self._best_ep = ep + self._best = cur + self.logger.info(f"Current metric: {cur} @ ep {ep}. Best metric: {self._best} @ ep {self._best_ep}") + + if ep - self._best_ep > self._patience: + self.workflow.early_stop = True + self.logger.info( + f"Validation metric has not been updated for {ep - self._best_ep} " + f"epochs (patience = {self._patience} epochs). Early stop.", + ) + + +class Checkpoint(Callback): + def __init__(self, path: str, interval: int) -> None: + super(Checkpoint, self).__init__() + + self._path = path + self._interval = interval + + def on_training_end(self, ep: int) -> None: + if ep % self._interval == 0: + self.training_manager.save(os.path.join(self._path, str(ep))) + self.logger.info(f"[Episode {ep}] All trainer states saved under {self._path}") + + +class MetricsRecorder(Callback): + def __init__(self, path: str) -> None: + super(MetricsRecorder, self).__init__() + + self._full_metrics: Dict[int, dict] = {} + self._valid_metrics: Dict[int, dict] = {} + self._path = path + + def _dump_metric_history(self) -> None: + if len(self._full_metrics) > 0: + metric_list = [self._full_metrics[ep] for ep in sorted(self._full_metrics.keys())] + df = pd.DataFrame.from_records(metric_list) + df.to_csv(os.path.join(self._path, "metrics_full.csv"), index=True) + if len(self._valid_metrics) > 0: + metric_list = [self._valid_metrics[ep] for ep in sorted(self._valid_metrics.keys())] + df = pd.DataFrame.from_records(metric_list) + df.to_csv(os.path.join(self._path, "metrics_valid.csv"), index=True) + + def on_training_end(self, ep: int) -> None: + if len(self.env_sampler.metrics) > 0: + metrics = copy.deepcopy(self.env_sampler.metrics) + metrics["ep"] = ep + if ep in self._full_metrics: + self._full_metrics[ep].update(metrics) + else: + self._full_metrics[ep] = metrics + self._dump_metric_history() + + def on_validation_end(self, ep: int) -> None: + if len(self.env_sampler.metrics) > 0: + metrics = copy.deepcopy(self.env_sampler.metrics) + metrics["ep"] = ep + if ep in self._full_metrics: + self._full_metrics[ep].update(metrics) + else: + self._full_metrics[ep] = metrics + if ep in self._valid_metrics: + self._valid_metrics[ep].update(metrics) + else: + self._valid_metrics[ep] = metrics + self._dump_metric_history() + + +class CallbackManager(object): + def __init__( + self, + workflow: TrainingWorkflow, + callbacks: List[Callback], + env_sampler: EnvSampler, + training_manager: TrainingManager, + logger: LoggerV2, + ) -> None: + super(CallbackManager, self).__init__() + + self._callbacks = callbacks + for callback in self._callbacks: + callback.workflow = workflow + callback.env_sampler = env_sampler + callback.training_manager = training_manager + callback.logger = logger + + def on_episode_start(self, ep: int) -> None: + for callback in self._callbacks: + callback.on_episode_start(ep) + + def on_episode_end(self, ep: int) -> None: + for callback in self._callbacks: + callback.on_episode_end(ep) + + def on_training_start(self, ep: int) -> None: + for callback in self._callbacks: + callback.on_training_start(ep) + + def on_training_end(self, ep: int) -> None: + for callback in self._callbacks: + callback.on_training_end(ep) + + def on_validation_start(self, ep: int) -> None: + for callback in self._callbacks: + callback.on_validation_start(ep) + + def on_validation_end(self, ep: int) -> None: + for callback in self._callbacks: + callback.on_validation_end(ep) + + def on_test_start(self, ep: int) -> None: + for callback in self._callbacks: + callback.on_test_start(ep) + + def on_test_end(self, ep: int) -> None: + for callback in self._callbacks: + callback.on_test_end(ep) diff --git a/maro/rl/workflows/config/parser.py b/maro/rl/workflows/config/parser.py index db52f065a..a94ec371f 100644 --- a/maro/rl/workflows/config/parser.py +++ b/maro/rl/workflows/config/parser.py @@ -76,6 +76,11 @@ def _validate_main_section(self) -> None: f"positive ints", ) + early_stop_patience = self._config["main"].get("early_stop_patience", None) + if early_stop_patience is not None: + if not isinstance(early_stop_patience, int) or early_stop_patience <= 0: + raise ValueError(f"Invalid early stop patience: {early_stop_patience}. Should be a positive integer.") + if "logging" in self._config["main"]: self._validate_logging_section("main", self._config["main"]["logging"]) @@ -196,9 +201,10 @@ def _validate_train_proxy_section(self, proxy_section: dict) -> None: raise TypeError(f"{self._validation_err_pfx}: 'training.proxy.backend' must be an int") def _validate_checkpointing_section(self, section: dict) -> None: - if "path" not in section: - raise KeyError(f"{self._validation_err_pfx}: missing field 'path' under section 'checkpointing'") - if not isinstance(section["path"], str): + ckpt_path = section.get("path", None) + if ckpt_path is None: + section["path"] = os.path.join(self._config["log_path"], "checkpoints") + elif not isinstance(section["path"], str): raise TypeError(f"{self._validation_err_pfx}: 'training.checkpointing.path' must be a string") if "interval" in section: @@ -231,10 +237,9 @@ def get_path_mapping(self, containerize: bool = False) -> dict: local/log/path -> "/logs" Defaults to False. """ - log_dir = os.path.dirname(self._config["log_path"]) path_map = { self._config["scenario_path"]: "/scenario" if containerize else self._config["scenario_path"], - log_dir: "/logs" if containerize else log_dir, + self._config["log_path"]: "/logs" if containerize else self._config["log_path"], } load_path = self._config["training"].get("load_path", None) @@ -286,12 +291,16 @@ def get_job_spec(self, containerize: bool = False) -> Dict[str, Tuple[str, Dict[ else: main_proc_env["EVAL_SCHEDULE"] = " ".join([str(val) for val in sorted(sch)]) + main_proc_env["NUM_EVAL_EPISODES"] = str(self._config["main"].get("num_eval_episodes", 1)) + if "early_stop_patience" in self._config["main"]: + main_proc_env["EARLY_STOP_PATIENCE"] = str(self._config["main"]["early_stop_patience"]) + load_path = self._config["training"].get("load_path", None) if load_path is not None: - env["main"]["LOAD_PATH"] = path_mapping[load_path] + main_proc_env["LOAD_PATH"] = path_mapping[load_path] load_episode = self._config["training"].get("load_episode", None) if load_episode is not None: - env["main"]["LOAD_EPISODE"] = str(load_episode) + main_proc_env["LOAD_EPISODE"] = str(load_episode) if "checkpointing" in self._config["training"]: conf = self._config["training"]["checkpointing"] @@ -385,9 +394,8 @@ def get_job_spec(self, containerize: bool = False) -> Dict[str, Tuple[str, Dict[ ) # All components write logs to the same file - log_dir, log_file = os.path.split(self._config["log_path"]) for _, vars in env.values(): - vars["LOG_PATH"] = os.path.join(path_mapping[log_dir], log_file) + vars["LOG_PATH"] = path_mapping[self._config["log_path"]] return env diff --git a/maro/rl/workflows/config/template.yml b/maro/rl/workflows/config/template.yml index 3464e9edc..ac514f78b 100644 --- a/maro/rl/workflows/config/template.yml +++ b/maro/rl/workflows/config/template.yml @@ -24,6 +24,8 @@ main: # A list indicates the episodes at the end of which policies are to be evaluated. Note that episode indexes are # 1-based. eval_schedule: 10 + early_stop_patience: 10 # Number of epochs waiting for a better validation metrics. Could be `null`. + num_eval_episodes: 10 # Number of Episodes to run in evaluation. # Minimum number of samples to start training in one epoch. The workflow will re-run experience collection # until we have at least `min_n_sample` of experiences. min_n_sample: 1 @@ -68,8 +70,9 @@ training: checkpointing: # Directory to save trainer snapshots under. Snapshot files created at different episodes will be saved under # separate folders named using episode numbers. For example, if a snapshot is created for a trainer named "dqn" - # at the end of episode 10, the file path would be "/path/to/your/checkpoint/folder/10/dqn.ckpt". - path: "/path/to/your/checkpoint/folder" + # at the end of episode 10, the file path would be "/path/to/your/checkpoint/folder/10/dqn.ckpt". If null, the + # default checkpoint folder would be created under `log_path`. + path: "/path/to/your/checkpoint/folder" # or `null` interval: 10 # Interval at which trained policies / models are persisted to disk. proxy: # Proxy settings. Ignored if training.mode is "simple". host: "127.0.0.1" # Proxy service host's IP address. Ignored if run in containerized environments. diff --git a/maro/rl/workflows/main.py b/maro/rl/workflows/main.py index 31de7caa1..e28a46035 100644 --- a/maro/rl/workflows/main.py +++ b/maro/rl/workflows/main.py @@ -14,21 +14,21 @@ from maro.rl.utils import get_torch_device from maro.rl.utils.common import float_or_none, get_env, int_or_none, list_or_none from maro.rl.utils.training import get_latest_ep -from maro.rl.workflows.utils import env_str_helper +from maro.rl.workflows.callback import CallbackManager, Checkpoint, EarlyStopping, MetricsRecorder from maro.utils import LoggerV2 class WorkflowEnvAttributes: def __init__(self) -> None: # Number of training episodes - self.num_episodes = int(env_str_helper(get_env("NUM_EPISODES"))) + self.num_episodes = int(get_env("NUM_EPISODES")) # Maximum number of steps in on round of sampling. self.num_steps = int_or_none(get_env("NUM_STEPS", required=False)) # Minimum number of data samples to start a round of training. If the data samples are insufficient, re-run # data sampling until we have at least `min_n_sample` data entries. - self.min_n_sample = int(env_str_helper(get_env("MIN_N_SAMPLE"))) + self.min_n_sample = int(get_env("MIN_N_SAMPLE")) # Path to store logs. self.log_path = get_env("LOG_PATH") @@ -46,6 +46,8 @@ def __init__(self) -> None: # Evaluating schedule. self.eval_schedule = list_or_none(get_env("EVAL_SCHEDULE", required=False)) + self.early_stop_patience = int_or_none(get_env("EARLY_STOP_PATIENCE", required=False)) + self.num_eval_episodes = int_or_none(get_env("NUM_EVAL_EPISODES", required=False)) # Restore configurations. self.load_path = get_env("LOAD_PATH", required=False) @@ -58,7 +60,7 @@ def __init__(self) -> None: # Parallel sampling configurations. self.parallel_rollout = self.env_sampling_parallelism is not None or self.env_eval_parallelism is not None if self.parallel_rollout: - self.port = int(env_str_helper(get_env("ROLLOUT_CONTROLLER_PORT"))) + self.port = int(get_env("ROLLOUT_CONTROLLER_PORT")) self.min_env_samples = int_or_none(get_env("MIN_ENV_SAMPLES", required=False)) self.grace_factor = float_or_none(get_env("GRACE_FACTOR", required=False)) @@ -67,13 +69,13 @@ def __init__(self) -> None: # Distributed training configurations. if self.train_mode != "simple": self.proxy_address = ( - env_str_helper(get_env("TRAIN_PROXY_HOST")), - int(env_str_helper(get_env("TRAIN_PROXY_FRONTEND_PORT"))), + str(get_env("TRAIN_PROXY_HOST")), + int(get_env("TRAIN_PROXY_FRONTEND_PORT")), ) self.logger = LoggerV2( "MAIN", - dump_path=self.log_path, + dump_path=os.path.join(self.log_path, "log.txt"), dump_mode="a", stdout_level=self.log_level_stdout, file_level=self.log_level_file, @@ -112,88 +114,111 @@ def main(rl_component_bundle: RLComponentBundle, env_attr: WorkflowEnvAttributes if args.evaluate_only: evaluate_only_workflow(rl_component_bundle, env_attr) else: - training_workflow(rl_component_bundle, env_attr) + TrainingWorkflow().run(rl_component_bundle, env_attr) -def training_workflow(rl_component_bundle: RLComponentBundle, env_attr: WorkflowEnvAttributes) -> None: - env_attr.logger.info("Start training workflow.") +class TrainingWorkflow(object): + def run(self, rl_component_bundle: RLComponentBundle, env_attr: WorkflowEnvAttributes) -> None: + env_attr.logger.info("Start training workflow.") - env_sampler = _get_env_sampler(rl_component_bundle, env_attr) + env_sampler = _get_env_sampler(rl_component_bundle, env_attr) - # evaluation schedule - env_attr.logger.info(f"Policy will be evaluated at the end of episodes {env_attr.eval_schedule}") - eval_point_index = 0 + # evaluation schedule + env_attr.logger.info(f"Policy will be evaluated at the end of episodes {env_attr.eval_schedule}") + eval_point_index = 0 - training_manager = TrainingManager( - rl_component_bundle=rl_component_bundle, - explicit_assign_device=(env_attr.train_mode == "simple"), - proxy_address=None if env_attr.train_mode == "simple" else env_attr.proxy_address, - logger=env_attr.logger, - ) + training_manager = TrainingManager( + rl_component_bundle=rl_component_bundle, + explicit_assign_device=(env_attr.train_mode == "simple"), + proxy_address=None if env_attr.train_mode == "simple" else env_attr.proxy_address, + logger=env_attr.logger, + ) - if env_attr.load_path: - assert isinstance(env_attr.load_path, str) + callbacks = [MetricsRecorder(path=env_attr.log_path)] + if env_attr.checkpoint_path is not None: + callbacks.append( + Checkpoint( + path=env_attr.checkpoint_path, + interval=1 if env_attr.checkpoint_interval is None else env_attr.checkpoint_interval, + ), + ) + if env_attr.early_stop_patience is not None: + callbacks.append(EarlyStopping(patience=env_attr.early_stop_patience)) + cbm = CallbackManager(self, callbacks, env_sampler, training_manager, env_attr.logger) + + if env_attr.load_path: + assert isinstance(env_attr.load_path, str) + + ep = env_attr.load_episode if env_attr.load_episode is not None else get_latest_ep(env_attr.load_path) + path = os.path.join(env_attr.load_path, str(ep)) + + loaded = env_sampler.load_policy_state(path) + env_attr.logger.info(f"Loaded policies {loaded} into env sampler from {path}") + + loaded = training_manager.load(path) + env_attr.logger.info(f"Loaded trainers {loaded} from {path}") + start_ep = ep + 1 + else: + start_ep = 1 + + # main loop + self.early_stop = False + for ep in range(start_ep, env_attr.num_episodes + 1): + if self.early_stop: # Might be set in `cbm.on_validation_end()` + break + + cbm.on_episode_start(ep) + + collect_time = training_time = 0.0 + total_experiences: List[List[ExpElement]] = [] + total_info_list: List[dict] = [] + n_sample = 0 + while n_sample < env_attr.min_n_sample: + tc0 = time.time() + result = env_sampler.sample( + policy_state=training_manager.get_policy_state() if not env_attr.is_single_thread else None, + num_steps=env_attr.num_steps, + ) + experiences: List[List[ExpElement]] = result["experiences"] + info_list: List[dict] = result["info"] + + n_sample += len(experiences[0]) + total_experiences.extend(experiences) + total_info_list.extend(info_list) + + collect_time += time.time() - tc0 + + env_sampler.post_collect(total_info_list, ep) + + tu0 = time.time() + env_attr.logger.info(f"Roll-out completed for episode {ep}. Training started...") + cbm.on_training_start(ep) + training_manager.record_experiences(total_experiences) + training_manager.train_step() + cbm.on_training_end(ep) + training_time += time.time() - tu0 + + # performance details + env_attr.logger.info( + f"ep {ep} - roll-out time: {collect_time:.2f} seconds, training time: {training_time:.2f} seconds", + ) + if env_attr.eval_schedule and ep == env_attr.eval_schedule[eval_point_index]: + cbm.on_validation_start(ep) - ep = env_attr.load_episode if env_attr.load_episode is not None else get_latest_ep(env_attr.load_path) - path = os.path.join(env_attr.load_path, str(ep)) + eval_point_index += 1 + result = env_sampler.eval( + policy_state=training_manager.get_policy_state() if not env_attr.is_single_thread else None, + num_episodes=env_attr.num_eval_episodes, + ) + env_sampler.post_evaluate(result["info"], ep) - loaded = env_sampler.load_policy_state(path) - env_attr.logger.info(f"Loaded policies {loaded} into env sampler from {path}") + cbm.on_validation_end(ep) - loaded = training_manager.load(path) - env_attr.logger.info(f"Loaded trainers {loaded} from {path}") - start_ep = ep + 1 - else: - start_ep = 1 - - # main loop - for ep in range(start_ep, env_attr.num_episodes + 1): - collect_time = training_time = 0.0 - total_experiences: List[List[ExpElement]] = [] - total_info_list: List[dict] = [] - n_sample = 0 - while n_sample < env_attr.min_n_sample: - tc0 = time.time() - result = env_sampler.sample( - policy_state=training_manager.get_policy_state() if not env_attr.is_single_thread else None, - num_steps=env_attr.num_steps, - ) - experiences: List[List[ExpElement]] = result["experiences"] - info_list: List[dict] = result["info"] - - n_sample += len(experiences[0]) - total_experiences.extend(experiences) - total_info_list.extend(info_list) - - collect_time += time.time() - tc0 - - env_sampler.post_collect(total_info_list, ep) - - env_attr.logger.info(f"Roll-out completed for episode {ep}. Training started...") - tu0 = time.time() - training_manager.record_experiences(total_experiences) - training_manager.train_step() - if env_attr.checkpoint_path and (not env_attr.checkpoint_interval or ep % env_attr.checkpoint_interval == 0): - assert isinstance(env_attr.checkpoint_path, str) - pth = os.path.join(env_attr.checkpoint_path, str(ep)) - training_manager.save(pth) - env_attr.logger.info(f"All trainer states saved under {pth}") - training_time += time.time() - tu0 - - # performance details - env_attr.logger.info( - f"ep {ep} - roll-out time: {collect_time:.2f} seconds, training time: {training_time:.2f} seconds", - ) - if env_attr.eval_schedule and ep == env_attr.eval_schedule[eval_point_index]: - eval_point_index += 1 - result = env_sampler.eval( - policy_state=training_manager.get_policy_state() if not env_attr.is_single_thread else None, - ) - env_sampler.post_evaluate(result["info"], ep) + cbm.on_episode_end(ep) - if isinstance(env_sampler, BatchEnvSampler): - env_sampler.exit() - training_manager.exit() + if isinstance(env_sampler, BatchEnvSampler): + env_sampler.exit() + training_manager.exit() def evaluate_only_workflow(rl_component_bundle: RLComponentBundle, env_attr: WorkflowEnvAttributes) -> None: @@ -210,7 +235,7 @@ def evaluate_only_workflow(rl_component_bundle: RLComponentBundle, env_attr: Wor loaded = env_sampler.load_policy_state(path) env_attr.logger.info(f"Loaded policies {loaded} into env sampler from {path}") - result = env_sampler.eval() + result = env_sampler.eval(num_episodes=env_attr.num_eval_episodes) env_sampler.post_evaluate(result["info"], -1) if isinstance(env_sampler, BatchEnvSampler): @@ -218,7 +243,7 @@ def evaluate_only_workflow(rl_component_bundle: RLComponentBundle, env_attr: Wor if __name__ == "__main__": - scenario_path = env_str_helper(get_env("SCENARIO_PATH")) + scenario_path = get_env("SCENARIO_PATH") scenario_path = os.path.normpath(scenario_path) sys.path.insert(0, os.path.dirname(scenario_path)) module = importlib.import_module(os.path.basename(scenario_path)) diff --git a/maro/rl/workflows/rollout_worker.py b/maro/rl/workflows/rollout_worker.py index 8343873b3..59cfa7b0d 100644 --- a/maro/rl/workflows/rollout_worker.py +++ b/maro/rl/workflows/rollout_worker.py @@ -8,21 +8,20 @@ from maro.rl.rl_component.rl_component_bundle import RLComponentBundle from maro.rl.rollout import RolloutWorker from maro.rl.utils.common import get_env, int_or_none -from maro.rl.workflows.utils import env_str_helper from maro.utils import LoggerV2 if __name__ == "__main__": - scenario_path = env_str_helper(get_env("SCENARIO_PATH")) + scenario_path = get_env("SCENARIO_PATH") scenario_path = os.path.normpath(scenario_path) sys.path.insert(0, os.path.dirname(scenario_path)) module = importlib.import_module(os.path.basename(scenario_path)) rl_component_bundle: RLComponentBundle = getattr(module, "rl_component_bundle") - worker_idx = int(env_str_helper(get_env("ID"))) + worker_idx = int(get_env("ID")) logger = LoggerV2( f"ROLLOUT-WORKER.{worker_idx}", - dump_path=get_env("LOG_PATH"), + dump_path=os.path.join(get_env("LOG_PATH"), f"ROLLOUT-WORKER.{worker_idx}.txt"), dump_mode="a", stdout_level=get_env("LOG_LEVEL_STDOUT", required=False, default="CRITICAL"), file_level=get_env("LOG_LEVEL_FILE", required=False, default="CRITICAL"), @@ -30,7 +29,7 @@ worker = RolloutWorker( idx=worker_idx, rl_component_bundle=rl_component_bundle, - producer_host=env_str_helper(get_env("ROLLOUT_CONTROLLER_HOST")), + producer_host=get_env("ROLLOUT_CONTROLLER_HOST"), producer_port=int_or_none(get_env("ROLLOUT_CONTROLLER_PORT")), logger=logger, ) diff --git a/maro/rl/workflows/train_worker.py b/maro/rl/workflows/train_worker.py index 4565c5b72..8fad5d5b6 100644 --- a/maro/rl/workflows/train_worker.py +++ b/maro/rl/workflows/train_worker.py @@ -8,11 +8,10 @@ from maro.rl.rl_component.rl_component_bundle import RLComponentBundle from maro.rl.training import TrainOpsWorker from maro.rl.utils.common import get_env, int_or_none -from maro.rl.workflows.utils import env_str_helper from maro.utils import LoggerV2 if __name__ == "__main__": - scenario_path = env_str_helper(get_env("SCENARIO_PATH")) + scenario_path = get_env("SCENARIO_PATH") scenario_path = os.path.normpath(scenario_path) sys.path.insert(0, os.path.dirname(scenario_path)) module = importlib.import_module(os.path.basename(scenario_path)) @@ -22,15 +21,15 @@ worker_idx = int_or_none(get_env("ID")) logger = LoggerV2( f"TRAIN-WORKER.{worker_idx}", - dump_path=get_env("LOG_PATH"), + dump_path=os.path.join(get_env("LOG_PATH"), f"TRAIN-WORKER.{worker_idx}.txt"), dump_mode="a", stdout_level=get_env("LOG_LEVEL_STDOUT", required=False, default="CRITICAL"), file_level=get_env("LOG_LEVEL_FILE", required=False, default="CRITICAL"), ) worker = TrainOpsWorker( - idx=int(env_str_helper(get_env("ID"))), + idx=int(get_env("ID")), rl_component_bundle=rl_component_bundle, - producer_host=env_str_helper(get_env("TRAIN_PROXY_HOST")), + producer_host=get_env("TRAIN_PROXY_HOST"), producer_port=int_or_none(get_env("TRAIN_PROXY_BACKEND_PORT")), logger=logger, ) diff --git a/maro/rl/workflows/utils.py b/maro/rl/workflows/utils.py deleted file mode 100644 index accfbe86f..000000000 --- a/maro/rl/workflows/utils.py +++ /dev/null @@ -1,9 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from typing import Optional - - -def env_str_helper(string: Optional[str]) -> str: - assert string is not None - return string diff --git a/tests/rl/algorithms/__init__.py b/tests/rl/algorithms/__init__.py deleted file mode 100644 index 9a0454564..000000000 --- a/tests/rl/algorithms/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. diff --git a/tests/rl/algorithms/ppo.py b/tests/rl/algorithms/ppo.py deleted file mode 100644 index 61b3c8576..000000000 --- a/tests/rl/algorithms/ppo.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from maro.rl.training.algorithms import PPOParams, PPOTrainer - -from .ac import MyVCriticNet, get_ac_policy - -get_ppo_policy = get_ac_policy - - -def get_ppo_trainer(name: str, state_dim: int) -> PPOTrainer: - return PPOTrainer( - name=name, - reward_discount=0.99, - params=PPOParams( - get_v_critic_net_func=lambda: MyVCriticNet(state_dim), - grad_iters=80, - lam=0.97, - clip_ratio=0.2, - ), - ) diff --git a/tests/rl/gym_wrapper/__init__.py b/tests/rl/gym_wrapper/__init__.py index 90be439f0..9a0454564 100644 --- a/tests/rl/gym_wrapper/__init__.py +++ b/tests/rl/gym_wrapper/__init__.py @@ -1,8 +1,2 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. - -from .rl_component_bundle import rl_component_bundle - -__all__ = [ - "rl_component_bundle", -] diff --git a/tests/rl/gym_wrapper/common.py b/tests/rl/gym_wrapper/common.py new file mode 100644 index 000000000..41287cd8f --- /dev/null +++ b/tests/rl/gym_wrapper/common.py @@ -0,0 +1,31 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from typing import cast + +from maro.simulator import Env +from maro.utils import set_seeds + +from tests.rl.gym_wrapper.simulator.business_engine import GymBusinessEngine + +set_seeds(123) + +env_conf = { + "topology": "Walker2d-v4", # HalfCheetah-v4, Hopper-v4, Walker2d-v4, Swimmer-v4, Ant-v4 + "start_tick": 0, + "durations": 100000, # Set a very large number + "options": { + "random_seed": None, + }, +} + +learn_env = Env(business_engine_cls=GymBusinessEngine, **env_conf) +test_env = Env(business_engine_cls=GymBusinessEngine, **env_conf) +num_agents = len(learn_env.agent_idx_list) + +gym_env = cast(GymBusinessEngine, learn_env.business_engine).gym_env +gym_action_space = gym_env.action_space +gym_state_dim = gym_env.observation_space.shape[0] +gym_action_dim = gym_action_space.shape[0] +action_lower_bound, action_upper_bound = gym_action_space.low, gym_action_space.high +action_limit = gym_action_space.high[0] diff --git a/tests/rl/gym_wrapper/config.py b/tests/rl/gym_wrapper/config.py deleted file mode 100644 index 0d37afcf4..000000000 --- a/tests/rl/gym_wrapper/config.py +++ /dev/null @@ -1,13 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -algorithm = "ppo" - -env_conf = { - "topology": "Walker2d-v4", - "start_tick": 0, - "durations": 5000, - "options": { - "random_seed": None, - }, -} diff --git a/tests/rl/gym_wrapper/env_sampler.py b/tests/rl/gym_wrapper/env_sampler.py index 0e1e3d30a..20d387b72 100644 --- a/tests/rl/gym_wrapper/env_sampler.py +++ b/tests/rl/gym_wrapper/env_sampler.py @@ -1,28 +1,45 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. -from typing import Any, Dict, Tuple, Union +from typing import Any, Dict, List, Tuple, Type, Union import numpy as np +from maro.rl.policy.abs_policy import AbsPolicy from maro.rl.rollout import AbsEnvSampler, CacheElement +from maro.rl.rollout.env_sampler import AbsAgentWrapper, SimpleAgentWrapper +from maro.simulator.core import Env from tests.rl.gym_wrapper.simulator.business_engine import GymBusinessEngine from tests.rl.gym_wrapper.simulator.common import Action, DecisionEvent -def _show_info(rewards: list, tag: str) -> None: - print( - f"[{tag}] Total N-steps = {sum([len(e) for e in rewards])}, " - f"N segments = {len(rewards)}, " - f"Average reward = {np.mean([sum(e) for e in rewards]):.4f}, " - f"Max reward = {np.max([sum(e) for e in rewards]):.4f}, " - f"Min reward = {np.min([sum(e) for e in rewards]):.4f}, " - f"Average N-steps = {np.mean([len(e) for e in rewards]):.1f}\n", - ) +class GymEnvSampler(AbsEnvSampler): + def __init__( + self, + learn_env: Env, + test_env: Env, + policies: List[AbsPolicy], + agent2policy: Dict[Any, str], + trainable_policies: List[str] = None, + agent_wrapper_cls: Type[AbsAgentWrapper] = SimpleAgentWrapper, + reward_eval_delay: int = None, + max_episode_length: int = None, + ) -> None: + super(GymEnvSampler, self).__init__( + learn_env=learn_env, + test_env=test_env, + policies=policies, + agent2policy=agent2policy, + trainable_policies=trainable_policies, + agent_wrapper_cls=agent_wrapper_cls, + reward_eval_delay=reward_eval_delay, + max_episode_length=max_episode_length, + ) + self._sample_rewards = [] + self._eval_rewards = [] -class GymEnvSampler(AbsEnvSampler): def _get_global_and_agent_state_impl( self, event: DecisionEvent, @@ -39,15 +56,38 @@ def _get_reward(self, env_action_dict: dict, event: Any, tick: int) -> Dict[Any, return {0: be.get_reward_at_tick(tick)} def _post_step(self, cache_element: CacheElement) -> None: - self._info["env_metric"] = self._env.metrics + if not (self._end_of_episode or self.truncated): + return + rewards = list(self._env.metrics["reward_record"].values()) + self._sample_rewards.append((len(rewards), np.sum(rewards))) def _post_eval_step(self, cache_element: CacheElement) -> None: - self._post_step(cache_element) + if not (self._end_of_episode or self.truncated): + return + rewards = list(self._env.metrics["reward_record"].values()) + self._eval_rewards.append((len(rewards), np.sum(rewards))) def post_collect(self, info_list: list, ep: int) -> None: - rewards = [list(e["env_metric"]["reward_record"].values()) for e in info_list] - _show_info(rewards, "Collect") + cur = { + "n_steps": sum([n for n, _ in self._sample_rewards]), + "n_segment": len(self._sample_rewards), + "avg_reward": np.mean([r for _, r in self._sample_rewards]), + "avg_n_steps": np.mean([n for n, _ in self._sample_rewards]), + "max_n_steps": np.max([n for n, _ in self._sample_rewards]), + "n_interactions": self._total_number_interactions, + } + self.metrics.update(cur) + # clear validation metrics + self.metrics = {k: v for k, v in self.metrics.items() if not k.startswith("val/")} + self._sample_rewards.clear() def post_evaluate(self, info_list: list, ep: int) -> None: - rewards = [list(e["env_metric"]["reward_record"].values()) for e in info_list] - _show_info(rewards, "Evaluate") + cur = { + "val/n_steps": sum([n for n, _ in self._eval_rewards]), + "val/n_segment": len(self._eval_rewards), + "val/avg_reward": np.mean([r for _, r in self._eval_rewards]), + "val/avg_n_steps": np.mean([n for n, _ in self._eval_rewards]), + "val/max_n_steps": np.max([n for n, _ in self._eval_rewards]), + } + self.metrics.update(cur) + self._eval_rewards.clear() diff --git a/tests/rl/gym_wrapper/rl_component_bundle.py b/tests/rl/gym_wrapper/rl_component_bundle.py deleted file mode 100644 index e19ce433f..000000000 --- a/tests/rl/gym_wrapper/rl_component_bundle.py +++ /dev/null @@ -1,70 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT license. - -from typing import cast - -from maro.rl.rl_component.rl_component_bundle import RLComponentBundle -from maro.simulator import Env - -from .config import algorithm, env_conf -from .env_sampler import GymEnvSampler -from tests.rl.gym_wrapper.simulator.business_engine import GymBusinessEngine - -learn_env = Env(business_engine_cls=GymBusinessEngine, **env_conf) -test_env = learn_env -num_agents = len(learn_env.agent_idx_list) - -gym_env = cast(GymBusinessEngine, learn_env.business_engine).gym_env -gym_state_dim = gym_env.observation_space.shape[0] -gym_action_dim = gym_env.action_space.shape[0] -action_lower_bound, action_upper_bound = gym_env.action_space.low, gym_env.action_space.high -action_limit = gym_env.action_space.high[0] - -agent2policy = {agent: f"{algorithm}_{agent}.policy" for agent in learn_env.agent_idx_list} - -if algorithm == "ac": - from tests.rl.algorithms.ac import get_ac_policy, get_ac_trainer - - policies = [ - get_ac_policy(f"{algorithm}_{i}.policy", action_lower_bound, action_upper_bound, gym_state_dim, gym_action_dim) - for i in range(num_agents) - ] - trainers = [get_ac_trainer(f"{algorithm}_{i}", gym_state_dim) for i in range(num_agents)] -elif algorithm == "ppo": - from tests.rl.algorithms.ppo import get_ppo_policy, get_ppo_trainer - - policies = [ - get_ppo_policy(f"{algorithm}_{i}.policy", action_lower_bound, action_upper_bound, gym_state_dim, gym_action_dim) - for i in range(num_agents) - ] - trainers = [get_ppo_trainer(f"{algorithm}_{i}", gym_state_dim) for i in range(num_agents)] -elif algorithm == "sac": - from tests.rl.algorithms.sac import get_sac_policy, get_sac_trainer - - policies = [ - get_sac_policy( - f"{algorithm}_{i}.policy", - action_lower_bound, - action_upper_bound, - gym_state_dim, - gym_action_dim, - action_limit, - ) - for i in range(num_agents) - ] - trainers = [get_sac_trainer(f"{algorithm}_{i}", gym_state_dim, gym_action_dim) for i in range(num_agents)] -else: - raise ValueError(f"Unsupported algorithm: {algorithm}") - - -rl_component_bundle = RLComponentBundle( - env_sampler=GymEnvSampler( - learn_env=learn_env, - test_env=test_env, - policies=policies, - agent2policy=agent2policy, - ), - agent2policy=agent2policy, - policies=policies, - trainers=trainers, -) diff --git a/tests/rl/gym_wrapper/simulator/business_engine.py b/tests/rl/gym_wrapper/simulator/business_engine.py index 626a0f6e2..d5b15153d 100644 --- a/tests/rl/gym_wrapper/simulator/business_engine.py +++ b/tests/rl/gym_wrapper/simulator/business_engine.py @@ -4,6 +4,7 @@ from typing import List, Optional, cast import gym +import numpy as np from maro.backends.frame import FrameBase, SnapshotList from maro.event_buffer import CascadeEvent, EventBuffer, MaroEvents @@ -36,7 +37,6 @@ def __init__( self._gym_scenario_name = topology self._gym_env = gym.make(self._gym_scenario_name) - self._seed = additional_options.get("random_seed", None) self.reset() @@ -81,7 +81,7 @@ def get_info_at_tick(self, tick: int) -> object: # TODO return self._info_record[tick] def reset(self, keep_seed: bool = False) -> None: - self._last_obs = self._gym_env.reset()[0] + self._last_obs = self._gym_env.reset(seed=np.random.randint(low=0, high=4096))[0] self._is_done = False self._truncated = False self._reward_record = {} diff --git a/tests/rl/log/Ant_1.png b/tests/rl/log/Ant_1.png new file mode 100644 index 000000000..f4f451d71 Binary files /dev/null and b/tests/rl/log/Ant_1.png differ diff --git a/tests/rl/log/Ant_11.png b/tests/rl/log/Ant_11.png new file mode 100644 index 000000000..930a258f9 Binary files /dev/null and b/tests/rl/log/Ant_11.png differ diff --git a/tests/rl/log/HalfCheetah_1.png b/tests/rl/log/HalfCheetah_1.png new file mode 100644 index 000000000..be1582bfe Binary files /dev/null and b/tests/rl/log/HalfCheetah_1.png differ diff --git a/tests/rl/log/HalfCheetah_11.png b/tests/rl/log/HalfCheetah_11.png new file mode 100644 index 000000000..2a94d104b Binary files /dev/null and b/tests/rl/log/HalfCheetah_11.png differ diff --git a/tests/rl/log/Hopper_1.png b/tests/rl/log/Hopper_1.png new file mode 100644 index 000000000..eeb54696f Binary files /dev/null and b/tests/rl/log/Hopper_1.png differ diff --git a/tests/rl/log/Hopper_11.png b/tests/rl/log/Hopper_11.png new file mode 100644 index 000000000..a3d576122 Binary files /dev/null and b/tests/rl/log/Hopper_11.png differ diff --git a/tests/rl/log/Swimmer_1.png b/tests/rl/log/Swimmer_1.png new file mode 100644 index 000000000..d69962dcd Binary files /dev/null and b/tests/rl/log/Swimmer_1.png differ diff --git a/tests/rl/log/Swimmer_11.png b/tests/rl/log/Swimmer_11.png new file mode 100644 index 000000000..c65db7d21 Binary files /dev/null and b/tests/rl/log/Swimmer_11.png differ diff --git a/tests/rl/log/Walker2d_1.png b/tests/rl/log/Walker2d_1.png new file mode 100644 index 000000000..c2aae720e Binary files /dev/null and b/tests/rl/log/Walker2d_1.png differ diff --git a/tests/rl/log/Walker2d_11.png b/tests/rl/log/Walker2d_11.png new file mode 100644 index 000000000..552460a8b Binary files /dev/null and b/tests/rl/log/Walker2d_11.png differ diff --git a/tests/rl/performance.md b/tests/rl/performance.md index 43849b553..9b7afdd8c 100644 --- a/tests/rl/performance.md +++ b/tests/rl/performance.md @@ -5,26 +5,37 @@ Some are compared to the benchmarks in [OpenAI Spinning Up](https://spinningup.o Limited by the environment version difference, there may be some gaps between the performance here and that in Spinning Up benchmarks. +## Experimental Setting + The hyper-parameters are set to align with those used in [Spinning Up](https://spinningup.openai.com/en/latest/spinningup/bench.html#experiment-details): -- Network of on-policy algorithms: size (64, 32) with tanh units for both policy and value function; -- Network of off-policy algorithms: size (256, 256) with relu units; -- Batch size for on-policy algorithms: 4000 steps of interaction per batch update; -- Batch size for off-policy algorithms: size 100 for each gradient descent step; +**Batch Size**: + +- For on-policy algorithms: 4000 steps of interaction per batch update; +- For off-policy algorithms: size 100 for each gradient descent step; + +**Network**: + +- For on-policy algorithms: size (64, 32) with tanh units for both policy and value function; +- For off-policy algorithms: size (256, 256) with relu units; + +**Performance metric**: -## Walker2d +- For on-policy algorithms: measured as the average trajectory return across the batch collected at each epoch; +- For off-policy algorithms: measured once every 10,000 steps by running the deterministic policy (or, in the case of SAC, the mean policy) without action noise for ten trajectories, and reporting the average return over those test trajectories; -### Benchmark in Spinning Up - PyTorch Version +**Total timesteps**: set to 3M for all task suites and algorithms. -- Environment version: Walker2d-v3 -- 3M timesteps +Other parameters are set to the values in *tests/rl/tasks/*. -![Walker2d: PyTorch Version](https://spinningup.openai.com/en/latest/_images/pytorch_walker2d_performance.svg) +## Performance Comparison -### Performance with MARO RL Toolkit +Five environments from the MuJoCo Gym task suite are reported in Spinning Up, they are: HalfCheetah, Hopper, Walker2d, Swimmer, and Ant. -- Environment version: Walker2d-v4 -- Training Mode: simple -- Rollout Mode: single -- Environment duration: 5000 ticks -- Num of episodes: 600 +| **Env** | **Spinning Up** | **MARO RL w/o Smooth** | **MARO RL w/ Smooth** | +|:---------------:|:---------------:|:----------------------:|:---------------------:| +| [**HalfCheetah**](https://gymnasium.farama.org/environments/mujoco/half_cheetah/) | ![Hab](https://spinningup.openai.com/en/latest/_images/pytorch_halfcheetah_performance.svg) | ![Ha1](./log/HalfCheetah_1.png) | ![Ha11](./log/HalfCheetah_11.png) | +| [**Hopper**](https://gymnasium.farama.org/environments/mujoco/hopper/) | ![Hob](https://spinningup.openai.com/en/latest/_images/pytorch_hopper_performance.svg) | ![Ho1](./log/Hopper_1.png) | ![Ho11](./log/Hopper_11.png) | +| [**Walker2d**](https://gymnasium.farama.org/environments/mujoco/walker2d/) | ![Wab](https://spinningup.openai.com/en/latest/_images/pytorch_walker2d_performance.svg) | ![Wa1](./log/Walker2d_1.png) | ![Wa11](./log/Walker2d_11.png) | +| [**Swimmer**](https://gymnasium.farama.org/environments/mujoco/swimmer/) | ![Swb](https://spinningup.openai.com/en/latest/_images/pytorch_swimmer_performance.svg) | ![Sw1](./log/Swimmer_1.png) | ![Sw11](./log/Swimmer_11.png) | +| [**Ant**](https://gymnasium.farama.org/environments/mujoco/ant/) | ![Anb](https://spinningup.openai.com/en/latest/_images/pytorch_ant_performance.svg) | ![An1](./log/Ant_1.png) | ![An11](./log/Ant_11.png) | diff --git a/tests/rl/plot.py b/tests/rl/plot.py new file mode 100644 index 000000000..664126d20 --- /dev/null +++ b/tests/rl/plot.py @@ -0,0 +1,90 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import argparse +import os +from typing import List, Tuple + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +LOG_DIR = "tests/rl/log" + +color_map = { + "ppo": "green", + "sac": "goldenrod", +} + + +def smooth(data: np.ndarray, window_size: int) -> np.ndarray: + if window_size > 1: + """ + smooth data with moving window average. + that is, + smoothed_y[t] = average(y[t-k], y[t-k+1], ..., y[t+k-1], y[t+k]) + where the "smooth" param is width of that window (2k+1) + """ + y = np.ones(window_size) + x = np.asarray(data) + z = np.ones_like(x) + smoothed_x = np.convolve(x, y, "same") / np.convolve(z, y, "same") + return smoothed_x + else: + return data + + +def get_off_policy_data(log_dir: str) -> Tuple[np.ndarray, np.ndarray]: + file_path = os.path.join(log_dir, "metrics_full.csv") + df = pd.read_csv(file_path) + x, y = df["n_interactions"], df["val/avg_reward"] + mask = ~np.isnan(y) + x, y = x[mask], y[mask] + return x, y + + +def get_on_policy_data(log_dir: str) -> Tuple[np.ndarray, np.ndarray]: + file_path = os.path.join(log_dir, "metrics_full.csv") + df = pd.read_csv(file_path) + x, y = df["n_interactions"], df["avg_reward"] + return x, y + + +def plot_performance_curves(title: str, dir_names: List[str], smooth_window_size: int) -> None: + for name in dir_names: + log_dir = os.path.join(LOG_DIR, name) + if not os.path.exists(log_dir): + continue + + if "ppo" in name: + algorithm = "ppo" + func = get_on_policy_data + elif "sac" in name: + algorithm = "sac" + func = get_off_policy_data + else: + raise "unknown algorithm name" + + x, y = func(log_dir) + y = smooth(y, smooth_window_size) + plt.plot(x, y, label=algorithm, color=color_map[algorithm]) + + plt.legend() + plt.title(title) + plt.xlabel("Total Env Interactions") + plt.ylabel(f"Average Trajectory Return (moving average with window size = {smooth_window_size})") + plt.savefig(os.path.join(LOG_DIR, f"{title}_{smooth_window_size}.png")) + plt.close() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--smooth", "-s", type=int, default=11, help="smooth window size") + args = parser.parse_args() + + for env_name in ["HalfCheetah", "Hopper", "Walker2d", "Swimmer", "Ant"]: + plot_performance_curves( + title=env_name, + dir_names=[f"{algorithm}_{env_name.lower()}" for algorithm in ["ppo", "sac"]], + smooth_window_size=args.smooth, + ) diff --git a/tests/rl/algorithms/ac.py b/tests/rl/tasks/ac/__init__.py similarity index 72% rename from tests/rl/algorithms/ac.py rename to tests/rl/tasks/ac/__init__.py index eea32fbcb..24cc961fc 100644 --- a/tests/rl/algorithms/ac.py +++ b/tests/rl/tasks/ac/__init__.py @@ -11,14 +11,26 @@ from maro.rl.model import ContinuousACBasedNet, VNet from maro.rl.model.fc_block import FullyConnected from maro.rl.policy import ContinuousRLPolicy +from maro.rl.rl_component.rl_component_bundle import RLComponentBundle from maro.rl.training.algorithms import ActorCriticParams, ActorCriticTrainer +from tests.rl.gym_wrapper.common import ( + action_lower_bound, + action_upper_bound, + gym_action_dim, + gym_state_dim, + learn_env, + num_agents, + test_env, +) +from tests.rl.gym_wrapper.env_sampler import GymEnvSampler + actor_net_conf = { - "hidden_dims": [64, 64], + "hidden_dims": [64, 32], "activation": torch.nn.Tanh, } critic_net_conf = { - "hidden_dims": [64, 64], + "hidden_dims": [64, 32], "activation": torch.nn.Tanh, } actor_learning_rate = 3e-4 @@ -95,3 +107,32 @@ def get_ac_trainer(name: str, state_dim: int) -> ActorCriticTrainer: lam=0.97, ), ) + + +algorithm = "ac" +agent2policy = {agent: f"{algorithm}_{agent}.policy" for agent in learn_env.agent_idx_list} +policies = [ + get_ac_policy(f"{algorithm}_{i}.policy", action_lower_bound, action_upper_bound, gym_state_dim, gym_action_dim) + for i in range(num_agents) +] +trainers = [get_ac_trainer(f"{algorithm}_{i}", gym_state_dim) for i in range(num_agents)] + +device_mapping = None +if torch.cuda.is_available(): + device_mapping = {f"{algorithm}_{i}.policy": "cuda:0" for i in range(num_agents)} + + +rl_component_bundle = RLComponentBundle( + env_sampler=GymEnvSampler( + learn_env=learn_env, + test_env=test_env, + policies=policies, + agent2policy=agent2policy, + ), + agent2policy=agent2policy, + policies=policies, + trainers=trainers, + device_mapping=device_mapping, +) + +__all__ = ["rl_component_bundle"] diff --git a/tests/rl/tasks/ac/config.yml b/tests/rl/tasks/ac/config.yml new file mode 100644 index 000000000..16ea56645 --- /dev/null +++ b/tests/rl/tasks/ac/config.yml @@ -0,0 +1,32 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +# Example RL config file for GYM scenario. +# Please refer to `maro/rl/workflows/config/template.yml` for the complete template and detailed explanations. + +job: gym_rl_workflow +scenario_path: "tests/rl/tasks/ac" +log_path: "tests/rl/log/ac" +main: + num_episodes: 1000 + num_steps: null + eval_schedule: 5 + num_eval_episodes: 10 + min_n_sample: 5000 + logging: + stdout: INFO + file: DEBUG +rollout: + logging: + stdout: INFO + file: DEBUG +training: + mode: simple + load_path: null + load_episode: null + checkpointing: + path: null + interval: 5 + logging: + stdout: INFO + file: DEBUG diff --git a/tests/rl/tasks/ddpg/__init__.py b/tests/rl/tasks/ddpg/__init__.py new file mode 100644 index 000000000..a36238f83 --- /dev/null +++ b/tests/rl/tasks/ddpg/__init__.py @@ -0,0 +1,140 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import torch +from torch.optim import Adam + +from maro.rl.model import QNet +from maro.rl.model.algorithm_nets.ddpg import ContinuousDDPGNet +from maro.rl.model.fc_block import FullyConnected +from maro.rl.policy import ContinuousRLPolicy +from maro.rl.rl_component.rl_component_bundle import RLComponentBundle +from maro.rl.training.algorithms import DDPGParams, DDPGTrainer + +from tests.rl.gym_wrapper.common import ( + action_limit, + action_lower_bound, + action_upper_bound, + gym_action_dim, + gym_state_dim, + learn_env, + num_agents, + test_env, +) +from tests.rl.gym_wrapper.env_sampler import GymEnvSampler + +actor_net_conf = { + "hidden_dims": [256, 256], + "activation": torch.nn.ReLU, + "output_activation": torch.nn.Tanh, +} +critic_net_conf = { + "hidden_dims": [256, 256], + "activation": torch.nn.Tanh, +} +actor_learning_rate = 1e-3 +critic_learning_rate = 1e-3 + + +class MyContinuousDDPGNet(ContinuousDDPGNet): + def __init__(self, state_dim: int, action_dim: int, action_limit: float) -> None: + super(MyContinuousDDPGNet, self).__init__(state_dim=state_dim, action_dim=action_dim) + + self._net = FullyConnected( + input_dim=state_dim, + output_dim=action_dim, + hidden_dims=actor_net_conf["hidden_dims"], + activation=actor_net_conf["activation"], + output_activation=actor_net_conf["output_activation"], + ) + self._optim = Adam(self._net.parameters(), lr=critic_learning_rate) + self._action_limit = action_limit + self._noise_scale = 0.1 # TODO + + def _get_actions_impl(self, states: torch.Tensor, exploring: bool) -> torch.Tensor: + action = self._net(states) * self._action_limit + if exploring: + action += torch.randn(self.action_dim) * self._noise_scale + action = torch.clamp(action, -self._action_limit, self._action_limit) + return action + + +class MyQCriticNet(QNet): + def __init__(self, state_dim: int, action_dim: int) -> None: + super(MyQCriticNet, self).__init__(state_dim=state_dim, action_dim=action_dim) + self._critic = FullyConnected( + input_dim=state_dim + action_dim, + output_dim=1, + hidden_dims=critic_net_conf["hidden_dims"], + activation=critic_net_conf["activation"], + ) + self._optim = Adam(self._critic.parameters(), lr=critic_learning_rate) + + def _get_q_values(self, states: torch.Tensor, actions: torch.Tensor) -> torch.Tensor: + return self._critic(torch.cat([states, actions], dim=1).float()).squeeze(-1) + + +def get_ddpg_policy( + name: str, + action_lower_bound: list, + action_upper_bound: list, + gym_state_dim: int, + gym_action_dim: int, + action_limit: float, +) -> ContinuousRLPolicy: + return ContinuousRLPolicy( + name=name, + action_range=(action_lower_bound, action_upper_bound), + policy_net=MyContinuousDDPGNet(gym_state_dim, gym_action_dim, action_limit), + ) + + +def get_ddpg_trainer(name: str, state_dim: int, action_dim: int) -> DDPGTrainer: + return DDPGTrainer( + name=name, + reward_discount=0.99, + replay_memory_capacity=1000000, + batch_size=100, + params=DDPGParams( + get_q_critic_net_func=lambda: MyQCriticNet(state_dim, action_dim), + num_epochs=20, + n_start_train=1000, + soft_update_coef=0.005, + ), + ) + + +algorithm = "ddpg" +agent2policy = {agent: f"{algorithm}_{agent}.policy" for agent in learn_env.agent_idx_list} +policies = [ + get_ddpg_policy( + f"{algorithm}_{i}.policy", + action_lower_bound, + action_upper_bound, + gym_state_dim, + gym_action_dim, + action_limit, + ) + for i in range(num_agents) +] +trainers = [get_ddpg_trainer(f"{algorithm}_{i}", gym_state_dim, gym_action_dim) for i in range(num_agents)] + +device_mapping = None +if torch.cuda.is_available(): + device_mapping = {f"{algorithm}_{i}.policy": "cuda:0" for i in range(num_agents)} + + +rl_component_bundle = RLComponentBundle( + env_sampler=GymEnvSampler( + learn_env=learn_env, + test_env=test_env, + policies=policies, + agent2policy=agent2policy, + ), + agent2policy=agent2policy, + policies=policies, + trainers=trainers, + device_mapping=device_mapping, +) + +__all__ = ["rl_component_bundle"] diff --git a/tests/rl/config.yml b/tests/rl/tasks/ddpg/config.yml similarity index 74% rename from tests/rl/config.yml rename to tests/rl/tasks/ddpg/config.yml index 43644e91d..a7b645de5 100644 --- a/tests/rl/config.yml +++ b/tests/rl/tasks/ddpg/config.yml @@ -8,13 +8,14 @@ # - python tests/rl/run.py tests/rl/config.yml job: gym_rl_workflow -scenario_path: "tests/rl/gym_wrapper" -log_path: "tests/rl/log/gym.txt" +scenario_path: "tests/rl/tasks/ddpg" +log_path: "tests/rl/log/ddpg" main: - num_episodes: 1000 - num_steps: null - eval_schedule: 5 - min_n_sample: 5000 + num_episodes: 25000 + num_steps: 200 + eval_schedule: 25 + num_eval_episodes: 10 + min_n_sample: 1 logging: stdout: INFO file: DEBUG @@ -27,8 +28,8 @@ training: load_path: null load_episode: null checkpointing: - path: "tests/rl/checkpoint/gym" - interval: 5 + path: null + interval: 25 logging: stdout: INFO file: DEBUG diff --git a/tests/rl/tasks/ppo/__init__.py b/tests/rl/tasks/ppo/__init__.py new file mode 100644 index 000000000..15fc71069 --- /dev/null +++ b/tests/rl/tasks/ppo/__init__.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import torch + +from maro.rl.rl_component.rl_component_bundle import RLComponentBundle +from maro.rl.training.algorithms.ppo import PPOParams, PPOTrainer + +from tests.rl.gym_wrapper.common import ( + action_lower_bound, + action_upper_bound, + gym_action_dim, + gym_state_dim, + learn_env, + num_agents, + test_env, +) +from tests.rl.gym_wrapper.env_sampler import GymEnvSampler +from tests.rl.tasks.ac import MyVCriticNet, get_ac_policy + +get_ppo_policy = get_ac_policy + + +def get_ppo_trainer(name: str, state_dim: int) -> PPOTrainer: + return PPOTrainer( + name=name, + reward_discount=0.99, + replay_memory_capacity=4000, + batch_size=4000, + params=PPOParams( + get_v_critic_net_func=lambda: MyVCriticNet(state_dim), + grad_iters=80, + lam=0.97, + clip_ratio=0.2, + ), + ) + + +algorithm = "ppo" +agent2policy = {agent: f"{algorithm}_{agent}.policy" for agent in learn_env.agent_idx_list} +policies = [ + get_ppo_policy(f"{algorithm}_{i}.policy", action_lower_bound, action_upper_bound, gym_state_dim, gym_action_dim) + for i in range(num_agents) +] +trainers = [get_ppo_trainer(f"{algorithm}_{i}", gym_state_dim) for i in range(num_agents)] + +device_mapping = None +if torch.cuda.is_available(): + device_mapping = {f"{algorithm}_{i}.policy": "cuda:0" for i in range(num_agents)} + +rl_component_bundle = RLComponentBundle( + env_sampler=GymEnvSampler( + learn_env=learn_env, + test_env=test_env, + policies=policies, + agent2policy=agent2policy, + max_episode_length=1000, + ), + agent2policy=agent2policy, + policies=policies, + trainers=trainers, + device_mapping=device_mapping, +) + +__all__ = ["rl_component_bundle"] diff --git a/tests/rl/tasks/ppo/config.yml b/tests/rl/tasks/ppo/config.yml new file mode 100644 index 000000000..de2412ee9 --- /dev/null +++ b/tests/rl/tasks/ppo/config.yml @@ -0,0 +1,32 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +# Example RL config file for GYM scenario. +# Please refer to `maro/rl/workflows/config/template.yml` for the complete template and detailed explanations. + +job: gym_rl_workflow +scenario_path: "tests/rl/tasks/ppo" +log_path: "tests/rl/log/ppo_walker2d" +main: + num_episodes: 1000 + num_steps: 4000 + eval_schedule: 5 + num_eval_episodes: 10 + min_n_sample: 1 + logging: + stdout: INFO + file: DEBUG +rollout: + logging: + stdout: INFO + file: DEBUG +training: + mode: simple + load_path: null + load_episode: null + checkpointing: + path: null + interval: 5 + logging: + stdout: INFO + file: DEBUG diff --git a/tests/rl/algorithms/sac.py b/tests/rl/tasks/sac/__init__.py similarity index 61% rename from tests/rl/algorithms/sac.py rename to tests/rl/tasks/sac/__init__.py index 828ec2470..1e033f12b 100644 --- a/tests/rl/algorithms/sac.py +++ b/tests/rl/tasks/sac/__init__.py @@ -6,23 +6,39 @@ import numpy as np import torch import torch.nn.functional as F +from gym import spaces from torch.distributions import Normal from torch.optim import Adam from maro.rl.model import ContinuousSACNet, QNet from maro.rl.model.fc_block import FullyConnected from maro.rl.policy import ContinuousRLPolicy +from maro.rl.rl_component.rl_component_bundle import RLComponentBundle from maro.rl.training.algorithms import SoftActorCriticParams, SoftActorCriticTrainer +from maro.rl.utils import ndarray_to_tensor + +from tests.rl.gym_wrapper.common import ( + action_limit, + action_lower_bound, + action_upper_bound, + gym_action_dim, + gym_action_space, + gym_state_dim, + learn_env, + num_agents, + test_env, +) +from tests.rl.gym_wrapper.env_sampler import GymEnvSampler actor_net_conf = { - "hidden_dims": [64, 64], - "activation": torch.nn.Tanh, + "hidden_dims": [256, 256], + "activation": torch.nn.ReLU, } critic_net_conf = { - "hidden_dims": [64, 64], - "activation": torch.nn.Tanh, + "hidden_dims": [256, 256], + "activation": torch.nn.ReLU, } -actor_learning_rate = 3e-4 +actor_learning_rate = 1e-3 critic_learning_rate = 1e-3 LOG_STD_MAX = 2 @@ -30,7 +46,7 @@ class MyContinuousSACNet(ContinuousSACNet): - def __init__(self, state_dim: int, action_dim: int, action_limit: float) -> None: + def __init__(self, state_dim: int, action_dim: int, action_limit: float, action_space: spaces.Space) -> None: super(MyContinuousSACNet, self).__init__(state_dim=state_dim, action_dim=action_dim) self._net = FullyConnected( @@ -45,6 +61,8 @@ def __init__(self, state_dim: int, action_dim: int, action_limit: float) -> None self._action_limit = action_limit self._optim = Adam(self.parameters(), lr=actor_learning_rate) + self._action_space = action_space + def _get_actions_with_logps_impl(self, states: torch.Tensor, exploring: bool) -> Tuple[torch.Tensor, torch.Tensor]: net_out = self._net(states.float()) mu = self._mu(net_out) @@ -61,6 +79,11 @@ def _get_actions_with_logps_impl(self, states: torch.Tensor, exploring: bool) -> return pi_action, logp_pi + def _get_random_actions_impl(self, states: torch.Tensor) -> torch.Tensor: + return torch.stack( + [ndarray_to_tensor(self._action_space.sample(), device=self._device) for _ in range(states.shape[0])], + ) + class MyQCriticNet(QNet): def __init__(self, state_dim: int, action_dim: int) -> None: @@ -88,7 +111,8 @@ def get_sac_policy( return ContinuousRLPolicy( name=name, action_range=(action_lower_bound, action_upper_bound), - policy_net=MyContinuousSACNet(gym_state_dim, gym_action_dim, action_limit), + policy_net=MyContinuousSACNet(gym_state_dim, gym_action_dim, action_limit, action_space=gym_action_space), + warmup=10000, ) @@ -96,9 +120,49 @@ def get_sac_trainer(name: str, state_dim: int, action_dim: int) -> SoftActorCrit return SoftActorCriticTrainer( name=name, reward_discount=0.99, + replay_memory_capacity=1000000, + batch_size=100, params=SoftActorCriticParams( get_q_critic_net_func=lambda: MyQCriticNet(state_dim, action_dim), - num_epochs=10, - n_start_train=10000, + update_target_every=1, + entropy_coef=0.2, + num_epochs=50, + n_start_train=1000, + soft_update_coef=0.005, ), ) + + +algorithm = "sac" +agent2policy = {agent: f"{algorithm}_{agent}.policy" for agent in learn_env.agent_idx_list} +policies = [ + get_sac_policy( + f"{algorithm}_{i}.policy", + action_lower_bound, + action_upper_bound, + gym_state_dim, + gym_action_dim, + action_limit, + ) + for i in range(num_agents) +] +trainers = [get_sac_trainer(f"{algorithm}_{i}", gym_state_dim, gym_action_dim) for i in range(num_agents)] + +device_mapping = None +if torch.cuda.is_available(): + device_mapping = {f"{algorithm}_{i}.policy": "cuda:0" for i in range(num_agents)} + +rl_component_bundle = RLComponentBundle( + env_sampler=GymEnvSampler( + learn_env=learn_env, + test_env=test_env, + policies=policies, + agent2policy=agent2policy, + ), + agent2policy=agent2policy, + policies=policies, + trainers=trainers, + device_mapping=device_mapping, +) + +__all__ = ["rl_component_bundle"] diff --git a/tests/rl/tasks/sac/config.yml b/tests/rl/tasks/sac/config.yml new file mode 100644 index 000000000..8e8fb7d7d --- /dev/null +++ b/tests/rl/tasks/sac/config.yml @@ -0,0 +1,32 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +# Example RL config file for GYM scenario. +# Please refer to `maro/rl/workflows/config/template.yml` for the complete template and detailed explanations. + +job: gym_rl_workflow +scenario_path: "tests/rl/tasks/sac" +log_path: "tests/rl/log/sac_walker2d" +main: + num_episodes: 80000 + num_steps: 50 + eval_schedule: 200 + num_eval_episodes: 10 + min_n_sample: 1 + logging: + stdout: INFO + file: DEBUG +rollout: + logging: + stdout: INFO + file: DEBUG +training: + mode: simple + load_path: null + load_episode: null + checkpointing: + path: null + interval: 200 + logging: + stdout: INFO + file: DEBUG