Refine RL workflow & tune RL models under GYM

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"]

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)

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

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

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"

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

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)

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

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

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

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,13 +140,16 @@ 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
  self._action_dim = action_dim
  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

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)