Distributed Framework

examples/distributed.py

@@ -0,0 +1,83 @@
+from genrl.distributed import (
+    Master,
+    ExperienceServer,
+    ParameterServer,
+    ActorNode,
+    LearnerNode,
+    WeightHolder,
+)
+from genrl.core import ReplayBuffer
+from genrl.agents import DDPG
+from genrl.trainers import DistributedTrainer
+import gym
+import argparse
+import torch.multiprocessing as mp
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-n", type=int)
+args = parser.parse_args()
+
+N_ACTORS = 2
+BUFFER_SIZE = 10
+MAX_ENV_STEPS = 100
+TRAIN_STEPS = 10
+BATCH_SIZE = 1
+
+
+def collect_experience(agent, experience_server_rref):
+    obs = agent.env.reset()
+    done = False
+    for i in range(MAX_ENV_STEPS):
+        action = agent.select_action(obs)
+        next_obs, reward, done, info = agent.env.step(action)
+        experience_server_rref.rpc_sync().push((obs, action, reward, next_obs, done))
+        obs = next_obs
+        if done:
+            break
+
+
+class MyTrainer(DistributedTrainer):
+    def __init__(self, agent, train_steps, batch_size):
+        super(MyTrainer, self).__init__(agent)
+        self.train_steps = train_steps
+        self.batch_size = batch_size
+
+    def train(self, parameter_server_rref, experience_server_rref):
+        i = 0
+        while i < self.train_steps:
+            batch = experience_server_rref.rpc_sync().sample(self.batch_size)
+            if batch is None:
+                continue
+            self.agent.update_params(batch, 1)
+            parameter_server_rref.rpc_sync().store_weights(self.agent.get_weights())
+            print(f"Trainer: {i + 1} / {self.train_steps} steps completed")
+            i += 1
+
+
+mp.set_start_method("fork")
+
+master = Master(world_size=6, address="localhost", port=29500)
+env = gym.make("Pendulum-v0")
+agent = DDPG("mlp", env)
+parameter_server = ParameterServer(
+    "param-0", master, WeightHolder(agent.get_weights()), rank=1
+)
+buffer = ReplayBuffer(BUFFER_SIZE)
+experience_server = ExperienceServer("experience-0", master, buffer, rank=2)
+trainer = MyTrainer(agent, TRAIN_STEPS, BATCH_SIZE)
+learner = LearnerNode(
+    "learner-0", master, parameter_server, experience_server, trainer, rank=3
+)
+actors = [
+    ActorNode(
+        f"actor-{i}",
+        master,
+        parameter_server,
+        experience_server,
+        learner,
+        agent,
+        collect_experience,
+        rank=i + 4,
+    )
+    for i in range(N_ACTORS)
+]

genrl/agents/deep/base/offpolicy.py

@@ -80,7 +80,7 @@ def _reshape_batch(self, batch: List):
         """
         return [*batch]
 
-    def sample_from_buffer(self, beta: float = None):
+    def sample_from_buffer(self, beta: float = None, batch = None):
         """Samples experiences from the buffer and converts them into usable formats
 
         Args:
@@ -89,11 +89,12 @@ def sample_from_buffer(self, beta: float = None):
         Returns:
             batch (:obj:`list`): Replay experiences sampled from the buffer
         """
-        # Samples from the buffer
-        if beta is not None:
-            batch = self.replay_buffer.sample(self.batch_size, beta=beta)
-        else:
-            batch = self.replay_buffer.sample(self.batch_size)
+        if batch is None:
+            # Samples from the buffer
+            if beta is not None:
+                batch = self.replay_buffer.sample(self.batch_size, beta=beta)
+            else:
+                batch = self.replay_buffer.sample(self.batch_size)
 
         states, actions, rewards, next_states, dones = self._reshape_batch(batch)
 
@@ -106,7 +107,7 @@ def sample_from_buffer(self, beta: float = None):
                 *[states, actions, rewards, next_states, dones, indices, weights]
             )
         else:
-            raise NotImplementedError
+            batch = ReplayBufferSamples(*[states, actions, rewards, next_states, dones])
         return batch
 
     def get_q_loss(self, batch: collections.namedtuple) -> torch.Tensor:
@@ -277,4 +278,4 @@ def load_weights(self, weights) -> None:
         Args:
             weights (:obj:`dict`): Dictionary of different neural net weights
         """
-        self.ac.load_state_dict(weights["weights"])
+        self.ac.load_state_dict(weights)

genrl/agents/deep/ddpg/ddpg.py

@@ -79,14 +79,14 @@ def _create_model(self) -> None:
         self.optimizer_policy = opt.Adam(self.ac.actor.parameters(), lr=self.lr_policy)
         self.optimizer_value = opt.Adam(self.ac.critic.parameters(), lr=self.lr_value)
 
-    def update_params(self, update_interval: int) -> None:
+    def update_params(self, batch, update_interval: int) -> None:
         """Update parameters of the model
 
         Args:
             update_interval (int): Interval between successive updates of the target model
         """
         for timestep in range(update_interval):
-            batch = self.sample_from_buffer()
+            batch = self.sample_from_buffer(batch=batch)
 
             value_loss = self.get_q_loss(batch)
             self.logs["value_loss"].append(value_loss.item())
@@ -123,6 +123,9 @@ def get_hyperparams(self) -> Dict[str, Any]:
         }
         return hyperparams
 
+    def get_weights(self):
+        return self.ac.state_dict()
+
     def get_logging_params(self) -> Dict[str, Any]:
         """Gets relevant parameters for logging
 

genrl/core/buffers.py

@@ -57,6 +57,9 @@ def sample(
                 :returns: (Tuple composing of `state`, `action`, `reward`,
         `next_state` and `done`)
         """
+        if batch_size > len(self.memory):
+            return None
+
         batch = random.sample(self.memory, batch_size)
         state, action, reward, next_state, done = map(np.stack, zip(*batch))
         return [

genrl/distributed/__init__.py

@@ -0,0 +1,5 @@
+from genrl.distributed.core import Master, Node
+from genrl.distributed.parameter_server import ParameterServer, WeightHolder
+from genrl.distributed.experience_server import ExperienceServer
+from genrl.distributed.actor import ActorNode
+from genrl.distributed.learner import LearnerNode

genrl/distributed/actor.py

@@ -0,0 +1,57 @@
+from genrl.distributed.core import Node
+from genrl.distributed.core import get_rref, store_rref
+import torch.distributed.rpc as rpc
+
+
+class ActorNode(Node):
+    def __init__(
+        self,
+        name,
+        master,
+        parameter_server,
+        experience_server,
+        learner,
+        agent,
+        collect_experience,
+        rank=None,
+    ):
+        super(ActorNode, self).__init__(name, master, rank)
+        self.parameter_server = parameter_server
+        self.experience_server = experience_server
+        self.init_proc(
+            target=self.act,
+            kwargs=dict(
+                parameter_server_name=parameter_server.name,
+                experience_server_name=experience_server.name,
+                learner_name=learner.name,
+                agent=agent,
+                collect_experience=collect_experience,
+            ),
+        )
+        self.start_proc()
+
+    @staticmethod
+    def act(
+        name,
+        world_size,
+        rank,
+        parameter_server_name,
+        experience_server_name,
+        learner_name,
+        agent,
+        collect_experience,
+        **kwargs,
+    ):
+        rpc.init_rpc(name=name, world_size=world_size, rank=rank)
+        print(f"{name}: RPC Initialised")
+        rref = rpc.RRef(agent)
+        store_rref(name, rref)
+        parameter_server_rref = get_rref(parameter_server_name)
+        experience_server_rref = get_rref(experience_server_name)
+        learner_rref = get_rref(learner_name)
+        print(f"{name}: Begining experience collection")
+        while not learner_rref.rpc_sync().is_done():
+            agent.load_weights(parameter_server_rref.rpc_sync().get_weights())
+            collect_experience(agent, experience_server_rref)
+
+        rpc.shutdown()

genrl/distributed/core.py

@@ -0,0 +1,148 @@
+import torch.distributed.rpc as rpc
+
+import threading
+
+from abc import ABC, abstractmethod
+import torch.multiprocessing as mp
+import os
+import time
+
+_rref_reg = {}
+_global_lock = threading.Lock()
+
+
+def _get_rref(idx):
+    global _rref_reg
+    with _global_lock:
+        if idx in _rref_reg.keys():
+            return _rref_reg[idx]
+        else:
+            return None
+
+
+def _store_rref(idx, rref):
+    global _rref_reg
+    with _global_lock:
+        if idx in _rref_reg.keys():
+            raise Warning(
+                f"Re-assigning RRef for key: {idx}. Make sure you are not using duplicate names for nodes"
+            )
+        _rref_reg[idx] = rref
+
+
+def get_rref(idx):
+    rref = rpc.rpc_sync("master", _get_rref, args=(idx,))
+    while rref is None:
+        time.sleep(0.5)
+        rref = rpc.rpc_sync("master", _get_rref, args=(idx,))
+    return rref
+
+
+def store_rref(idx, rref):
+    rpc.rpc_sync("master", _store_rref, args=(idx, rref))
+
+
+def set_environ(address, port):
+    os.environ["MASTER_ADDR"] = str(address)
+    os.environ["MASTER_PORT"] = str(port)
+
+
+class Node:
+    def __init__(self, name, master, rank):
+        self._name = name
+        self.master = master
+        if rank is None:
+            self._rank = master.node_count
+        elif rank >= 0 and rank < master.world_size:
+            self._rank = rank
+        elif rank >= master.world_size:
+            raise ValueError("Specified rank greater than allowed by world size")
+        else:
+            raise ValueError("Invalid value of rank")
+        self.p = None
+
+    def __del__(self):
+        if self.p is None:
+            raise RuntimeWarning(
+                "Removing node when process was not initialised properly"
+            )
+        else:
+            self.p.join()
+
+    @staticmethod
+    def _target_wrapper(target, **kwargs):
+        pid = os.getpid()
+        print(f"Starting {kwargs['name']} with pid {pid}")
+        set_environ(kwargs["master_address"], kwargs["master_port"])
+        target(**kwargs)
+        print(f"Shutdown {kwargs['name']} with pid {pid}")
+
+    def init_proc(self, target, kwargs):
+        kwargs.update(
+            dict(
+                name=self.name,
+                master_address=self.master.address,
+                master_port=self.master.port,
+                world_size=self.master.world_size,
+                rank=self.rank,
+            )
+        )
+        self.p = mp.Process(target=self._target_wrapper, args=(target,), kwargs=kwargs)
+
+    def start_proc(self):
+        if self.p is None:
+            raise RuntimeError("Trying to start uninitialised process")
+        self.p.start()
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def rref(self):
+        return get_rref(self.name)
+
+    @property
+    def rank(self):
+        return self._rank
+
+
+def _run_master(world_size):
+    print(f"Starting master at {os.getpid()}")
+    rpc.init_rpc("master", rank=0, world_size=world_size)
+    rpc.shutdown()
+
+
+class Master:
+    def __init__(self, world_size, address="localhost", port=29501):
+        set_environ(address, port)
+        self._world_size = world_size
+        self._address = address
+        self._port = port
+        self._node_counter = 0
+        self.p = mp.Process(target=_run_master, args=(world_size,))
+        self.p.start()
+
+    def __del__(self):
+        if self.p is None:
+            raise RuntimeWarning(
+                "Shutting down master when it was not initialised properly"
+            )
+        else:
+            self.p.join()
+
+    @property
+    def world_size(self):
+        return self._world_size
+
+    @property
+    def address(self):
+        return self._address
+
+    @property
+    def port(self):
+        return self._port
+
+    @property
+    def node_count(self):
+        return self._node_counter

genrl/distributed/experience_server.py

@@ -0,0 +1,23 @@
+from genrl.distributed import Node
+from genrl.distributed.core import store_rref
+
+import torch.distributed.rpc as rpc
+
+
+class ExperienceServer(Node):
+    def __init__(self, name, master, buffer, rank=None):
+        super(ExperienceServer, self).__init__(name, master, rank)
+        self.init_proc(
+            target=self.run_paramater_server,
+            kwargs=dict(buffer=buffer),
+        )
+        self.start_proc()
+
+    @staticmethod
+    def run_paramater_server(name, world_size, rank, buffer, **kwargs):
+        rpc.init_rpc(name=name, world_size=world_size, rank=rank)
+        print(f"{name}: Initialised RPC")
+        rref = rpc.RRef(buffer)
+        store_rref(name, rref)
+        print(f"{name}: Serving experience buffer")
+        rpc.shutdown()