| 117 | |
| 118 | |
| 119 | class Agent: |
| 120 | def __init__(self, world_size, batch=True): |
| 121 | self.ob_rrefs = [] |
| 122 | self.agent_rref = RRef(self) |
| 123 | self.rewards = {} |
| 124 | self.policy = Policy(batch).cuda() |
| 125 | self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2) |
| 126 | self.running_reward = 0 |
| 127 | |
| 128 | for ob_rank in range(1, world_size): |
| 129 | ob_info = rpc.get_worker_info(OBSERVER_NAME.format(ob_rank)) |
| 130 | self.ob_rrefs.append(remote(ob_info, Observer, args=(batch,))) |
| 131 | self.rewards[ob_info.id] = [] |
| 132 | |
| 133 | self.states = torch.zeros(len(self.ob_rrefs), 1, 4) |
| 134 | self.batch = batch |
| 135 | # With batching, saved_log_probs contains a list of tensors, where each |
| 136 | # tensor contains probs from all observers in one step. |
| 137 | # Without batching, saved_log_probs is a dictionary where the key is the |
| 138 | # observer id and the value is a list of probs for that observer. |
| 139 | self.saved_log_probs = [] if self.batch else {k:[] for k in range(len(self.ob_rrefs))} |
| 140 | self.future_actions = torch.futures.Future() |
| 141 | self.lock = threading.Lock() |
| 142 | self.pending_states = len(self.ob_rrefs) |
| 143 | |
| 144 | @staticmethod |
| 145 | @rpc.functions.async_execution |
| 146 | def select_action_batch(agent_rref, ob_id, state): |
| 147 | r""" |
| 148 | Batching select_action: In each step, the agent waits for states from |
| 149 | all observers, and process them together. This helps to reduce the |
| 150 | number of CUDA kernels launched and hence speed up amortized inference |
| 151 | speed. |
| 152 | """ |
| 153 | self = agent_rref.local_value() |
| 154 | self.states[ob_id].copy_(state) |
| 155 | future_action = self.future_actions.then( |
| 156 | lambda future_actions: future_actions.wait()[ob_id].item() |
| 157 | ) |
| 158 | |
| 159 | with self.lock: |
| 160 | self.pending_states -= 1 |
| 161 | if self.pending_states == 0: |
| 162 | self.pending_states = len(self.ob_rrefs) |
| 163 | probs = self.policy(self.states.cuda()) |
| 164 | m = Categorical(probs) |
| 165 | actions = m.sample() |
| 166 | self.saved_log_probs.append(m.log_prob(actions).t()[0]) |
| 167 | future_actions = self.future_actions |
| 168 | self.future_actions = torch.futures.Future() |
| 169 | future_actions.set_result(actions.cpu()) |
| 170 | return future_action |
| 171 | |
| 172 | @staticmethod |
| 173 | def select_action(agent_rref, ob_id, state): |
| 174 | r""" |
| 175 | Non-batching select_action, return the action right away. |
| 176 | """ |