[RLlib] Add Optimizer State To Learner get_state

rllib/BUILD

@@ -2283,8 +2283,8 @@ py_test(
 
 py_test(
  name = "utils/tests/test_torch_utils",
- tags = ["team:rllib", "utils"],
- size = "small",
+ tags = ["team:rllib", "utils", "gpu"],
+ size = "medium",
  srcs = ["utils/tests/test_torch_utils.py"]
 )
 

rllib/core/learner/learner.py

@@ -849,14 +849,29 @@ def set_state(self, state: Mapping[str, Any]) -> None:
 
  Args:
  state: The state of the optimizer and module. Can be obtained
- from `get_state`.
+ from `get_state`. State is a dictionary with two keys:
+ "module_state" and "optimizer_state". The value of each key
+ is a dictionary that can be passed to `set_weights` and
+ `set_optimizer_weights` respectively.
 
  """
  # TODO (Kourosh): We have both get(set)_state and get(set)_weights. I think
  # having both can become confusing. Can we simplify this API requirement?
  self._check_is_built()
  # TODO: once we figure out the optimizer format, we can set/get the state
- self._module.set_state(state.get("module_state", {}))
+ if "module_state" not in state:
+ raise ValueError(
+ "state must have a key 'module_state' for the module weights"
+ )
+ if "optimizer_state" not in state:
+ raise ValueError(
+ "state must have a key 'optimizer_state' for the optimizer weights"
+ )
+
+ module_state = state.get("module_state")
+ optimizer_state = state.get("optimizer_state")
+ self.set_weights(module_state)
+ self.set_optimizer_weights(optimizer_state)
 
  def get_state(self) -> Mapping[str, Any]:
  """Get the state of the learner.
@@ -867,7 +882,29 @@ def get_state(self) -> Mapping[str, Any]:
  """
  self._check_is_built()
  # TODO: once we figure out the optimizer format, we can set/get the state
- return {"module_state": self._module.get_state()}
+ return {
+ "module_state": self.get_weights(),
+ "optimizer_state": self.get_optimizer_weights(),
+ }
+ # return {"module_state": self.get_weights(), "optimizer_state": {}}
+
+ def set_optimizer_weights(self, weights: Mapping[str, Any]) -> None:
+ """Set the weights of the optimizer.
+
+ Args:
+ weights: The weights of the optimizer.
+
+ """
+ raise NotImplementedError
+
+ def get_optimizer_weights(self) -> Mapping[str, Any]:
+ """Get the weights of the optimizer.
+
+ Returns:
+ The weights of the optimizer.
+
+ """
+ raise NotImplementedError
 
  def _get_metadata(self) -> Dict[str, Any]:
  metadata = {

rllib/core/learner/tests/test_learner_group.py

@@ -36,7 +36,7 @@
 
 LOCAL_SCALING_CONFIGS = {
  "local-cpu": LearnerGroupScalingConfig(num_workers=0, num_gpus_per_worker=0),
- "local-gpu": LearnerGroupScalingConfig(num_workers=0, num_gpus_per_worker=0.5),
+ "local-gpu": LearnerGroupScalingConfig(num_workers=0, num_gpus_per_worker=1),
 }
 
 
@@ -45,6 +45,17 @@
 @ray.remote(num_gpus=1)
 class RemoteTrainingHelper:
  def local_training_helper(self, fw, scaling_mode) -> None:
+ if fw == "torch":
+ import torch
+
+ torch.manual_seed(0)
+ elif fw == "tf":
+ import tensorflow as tf
+
+ # this is done by rllib already inside of the policy class, but we need to
+ # do it here for testing purposes
+ tf.compat.v1.enable_eager_execution()
+ tf.random.set_seed(0)
  env = gym.make("CartPole-v1")
  scaling_config = LOCAL_SCALING_CONFIGS[scaling_mode]
  lr = 1e-3
@@ -71,13 +82,25 @@ def local_training_helper(self, fw, scaling_mode) -> None:
 
  # make the state of the learner and the local learner_group identical
  local_learner.set_state(learner_group.get_state())
+ # learner_group.set_state(learner_group.get_state())
+ check(local_learner.get_state(), learner_group.get_state())
 
  # do another update
  batch = reader.next()
  ma_batch = MultiAgentBatch(
  {new_module_id: batch, DEFAULT_POLICY_ID: batch}, env_steps=batch.count
  )
- check(local_learner.update(ma_batch), learner_group.update(ma_batch))
+ # the optimizer state is not initialized fully until the first time that
+ # training is completed. A call to get state before that won't contain the
+ # optimizer state. So we do a dummy update here to initialize the optimizer
+ local_learner.update(ma_batch)
+ learner_group.update(ma_batch)
+
+ check(local_learner.get_state(), learner_group.get_state())
+ local_learner_results = local_learner.update(ma_batch)
+ learner_group_results = learner_group.update(ma_batch)
+
+ check(local_learner_results, learner_group_results)
 
  check(local_learner.get_state(), learner_group.get_state())
 

rllib/core/learner/tf/tf_learner.py

@@ -270,6 +270,25 @@ def _load_optimizers(self, path: Union[str, pathlib.Path]) -> None:
  def set_weights(self, weights: Mapping[str, Any]) -> None:
  self._module.set_state(weights)
 
+ @override(Learner)
+ def get_optimizer_weights(self) -> Mapping[str, Any]:
+ optim_weights = {}
+ with tf.init_scope():
+ for name, optim in self._named_optimizers.items():
+ optim_weights[name] = [var.numpy() for var in optim.variables()]
+ return optim_weights
+
+ @override(Learner)
+ def set_optimizer_weights(self, weights: Mapping[str, Any]) -> None:
+ for name, weight_array in weights.items():
+ if name not in self._named_optimizers:
+ raise ValueError(
+ f"Optimizer {name} in weights is not known."
+ f"Known optimizers are {self._named_optimizers.keys()}"
+ )
+ optim = self._named_optimizers[name]
+ optim.set_weights(weight_array)
+
  @override(Learner)
  def get_param_ref(self, param: ParamType) -> Hashable:
  return param.ref()

rllib/core/learner/torch/torch_learner.py

@@ -27,9 +27,13 @@
 from ray.rllib.core.rl_module.torch.torch_rl_module import TorchDDPRLModule
 from ray.rllib.policy.sample_batch import MultiAgentBatch
 from ray.rllib.utils.annotations import override
-from ray.rllib.utils.torch_utils import clip_gradients, convert_to_torch_tensor
 from ray.rllib.utils.typing import TensorType
 from ray.rllib.utils.nested_dict import NestedDict
+from ray.rllib.utils.torch_utils import (
+ clip_gradients,
+ convert_to_torch_tensor,
+ copy_torch_tensors,
+)
 from ray.rllib.utils.framework import try_import_torch
 
 torch, nn = try_import_torch()
@@ -119,16 +123,42 @@ def set_weights(self, weights: Mapping[str, Any]) -> None:
  def _save_optimizers(self, path: Union[str, pathlib.Path]) -> None:
  path = pathlib.Path(path)
  path.mkdir(parents=True, exist_ok=True)
- for name, optim in self._named_optimizers.items():
- torch.save(optim.state_dict(), path / f"{name}.pt")
+ optim_weights = self.get_optimizer_weights()
+ for name, weights in optim_weights.items():
+ torch.save(weights, path / f"{name}.pt")
 
  @override(Learner)
  def _load_optimizers(self, path: Union[str, pathlib.Path]) -> None:
  path = pathlib.Path(path)
  if not path.exists():
  raise ValueError(f"Directory {path} does not exist.")
+ weights = {}
+ for name in self._named_optimizers.keys():
+ weights[name] = torch.load(path / f"{name}.pt")
+ self.set_optimizer_weights(weights)
+
+ @override(Learner)
+ def get_optimizer_weights(self) -> Mapping[str, Any]:
+ optimizer_name_weights = {}
  for name, optim in self._named_optimizers.items():
- optim.load_state_dict(torch.load(path / f"{name}.pt"))
+ optim_state_dict = optim.state_dict()
+ optim_state_dict_cpu = copy_torch_tensors(optim_state_dict, device="cpu")
+ optimizer_name_weights[name] = optim_state_dict_cpu
+ return optimizer_name_weights
+
+ @override(Learner)
+ def set_optimizer_weights(self, weights: Mapping[str, Any]) -> None:
+ for name, weight_dict in weights.items():
+ if name not in self._named_optimizers:
+ raise ValueError(
+ f"Optimizer {name} in weights is not known."
+ f"Known optimizers are {self._named_optimizers.keys()}"
+ )
+ optim = self._named_optimizers[name]
+ weight_dict_correct_device = copy_torch_tensors(
+ weight_dict, device=self._device
+ )
+ optim.load_state_dict(weight_dict_correct_device)
 
  @override(Learner)
  def get_param_ref(self, param: ParamType) -> Hashable:

rllib/utils/tests/test_torch_utils.py

@@ -4,7 +4,10 @@
 import torch.cuda
 
 import ray
-from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+from ray.rllib.utils.torch_utils import (
+ convert_to_torch_tensor,
+ copy_torch_tensors,
+)
 
 
 class TestTorchUtils(unittest.TestCase):
@@ -43,6 +46,54 @@ def test_convert_to_torch_tensor(self):
  self.assertTrue(converted["b"].dtype is torch.float32)
  self.assertTrue(converted["c"] is None)
 
+ def test_copy_torch_tensors(self):
+ array = np.array([1, 2, 3], dtype=np.float32)
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+ tensor = torch.from_numpy(array).to(device)
+ tensor_2 = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64).to(device)
+
+ # Test single tensor
+ copied_tensor = copy_torch_tensors(tensor, device)
+ self.assertTrue(copied_tensor.device == device)
+ self.assertNotEqual(id(copied_tensor), id(tensor))
+ self.assertTrue(all(copied_tensor == tensor))
+
+ # check that dtypes aren't modified
+ copied_tensor_2 = copy_torch_tensors(tensor_2, device)
+ self.assertTrue(copied_tensor_2.dtype == tensor_2.dtype)
+ self.assertFalse(copied_tensor_2.dtype == torch.float32)
+
+ # Test nested structure can be converted
+ nested_structure = {"a": tensor, "b": tensor_2, "c": 1}
+ copied_nested_structure = copy_torch_tensors(nested_structure, device)
+ self.assertTrue(copied_nested_structure["a"].device == device)
+ self.assertTrue(copied_nested_structure["b"].device == device)
+ self.assertTrue(copied_nested_structure["c"] == 1)
+ self.assertNotEqual(id(copied_nested_structure["a"]), id(tensor))
+ self.assertNotEqual(id(copied_nested_structure["b"]), id(tensor_2))
+ self.assertTrue(all(copied_nested_structure["a"] == tensor))
+ self.assertTrue(all(copied_nested_structure["b"] == tensor_2))
+
+ # if gpu is available test moving tensor from cpu to gpu and vice versa
+ if torch.cuda.is_available():
+ tensor = torch.from_numpy(array).to("cpu")
+ copied_tensor = copy_torch_tensors(tensor, "cuda:0")
+ self.assertFalse(copied_tensor.device == torch.device("cpu"))
+ self.assertTrue(copied_tensor.device == torch.device("cuda:0"))
+ self.assertNotEqual(id(copied_tensor), id(tensor))
+ self.assertTrue(
+ all(copied_tensor.detach().cpu().numpy() == tensor.detach().numpy())
+ )
+
+ tensor = torch.from_numpy(array).to("cuda:0")
+ copied_tensor = copy_torch_tensors(tensor, "cpu")
+ self.assertFalse(copied_tensor.device == torch.device("cuda:0"))
+ self.assertTrue(copied_tensor.device == torch.device("cpu"))
+ self.assertNotEqual(id(copied_tensor), id(tensor))
+ self.assertTrue(
+ all(copied_tensor.detach().numpy() == tensor.detach().cpu().numpy())
+ )
+
 
 if __name__ == "__main__":
  import pytest

rllib/utils/torch_utils.py

@@ -231,6 +231,37 @@ def mapping(item):
  return tree.map_structure(mapping, x)
 
 
+@PublicAPI
+def copy_torch_tensors(x: TensorStructType, device: Optional[str] = None):
+ """Creates a copy of `x` and makes deep copies torch.Tensors in x.
+
+ Also moves the copied tensors to the specified device (if not None).
+
+ Note if an object in x is not a torch.Tensor, it will be shallow-copied.
+
+ Args:
+ x : Any (possibly nested) struct possibly containing torch.Tensors.
+ device : The device to move the tensors to.
+
+ Returns:
+ Any: A new struct with the same structure as `x`, but with all
+ torch.Tensors deep-copied and moved to the specified device.
+
+ """
+
+ def mapping(item):
+ if isinstance(item, torch.Tensor):
+ return (
+ torch.clone(item.detach())
+ if device is None
+ else item.detach().to(device)
+ )
+ else:
+ return item
+
+ return tree.map_structure(mapping, x)
+
+
 @PublicAPI
 def explained_variance(y: TensorType, pred: TensorType) -> TensorType:
  """Computes the explained variance for a pair of labels and predictions.