def _compare_two_policies(policy1: TorchPolicy, policy2: TorchPolicy) -> None: """ Make sure two policies have the same output for the same input. """ policy1.actor = policy1.actor.to(default_device()) policy2.actor = policy2.actor.to(default_device()) decision_step, _ = mb.create_steps_from_behavior_spec( policy1.behavior_spec, num_agents=1) np_obs = decision_step.obs masks = policy1._extract_masks(decision_step) memories = torch.as_tensor( policy1.retrieve_memories(list(decision_step.agent_id))).unsqueeze(0) tensor_obs = [ModelUtils.list_to_tensor(obs) for obs in np_obs] with torch.no_grad(): _, log_probs1, _, _ = policy1.sample_actions(tensor_obs, masks=masks, memories=memories) _, log_probs2, _, _ = policy2.sample_actions(tensor_obs, masks=masks, memories=memories) np.testing.assert_array_equal( ModelUtils.to_numpy(log_probs1.all_discrete_tensor), ModelUtils.to_numpy(log_probs2.all_discrete_tensor), )
def evaluate(self, decision_requests: DecisionSteps, global_agent_ids: List[str]) -> Dict[str, Any]: """ Evaluates policy for the agent experiences provided. :param global_agent_ids: :param decision_requests: DecisionStep object containing inputs. :return: Outputs from network as defined by self.inference_dict. """ obs = decision_requests.obs masks = self._extract_masks(decision_requests) tensor_obs = [torch.as_tensor(np_ob) for np_ob in obs] memories = torch.as_tensor( self.retrieve_memories(global_agent_ids)).unsqueeze(0) run_out = {} with torch.no_grad(): action, log_probs, entropy, memories = self.sample_actions( tensor_obs, masks=masks, memories=memories) action_tuple = action.to_action_tuple() run_out["action"] = action_tuple # This is the clipped action which is not saved to the buffer # but is exclusively sent to the environment. env_action_tuple = action.to_action_tuple(clip=self._clip_action) run_out["env_action"] = env_action_tuple run_out["log_probs"] = log_probs.to_log_probs_tuple() run_out["entropy"] = ModelUtils.to_numpy(entropy) run_out["learning_rate"] = 0.0 if self.use_recurrent: run_out["memory_out"] = ModelUtils.to_numpy(memories).squeeze(0) return run_out
def get_trajectory_value_estimates( self, batch: AgentBuffer, next_obs: List[np.ndarray], done: bool) -> Tuple[Dict[str, np.ndarray], Dict[str, float]]: n_obs = len(self.policy.behavior_spec.observation_specs) current_obs = ObsUtil.from_buffer(batch, n_obs) # Convert to tensors current_obs = [ModelUtils.list_to_tensor(obs) for obs in current_obs] next_obs = [ModelUtils.list_to_tensor(obs) for obs in next_obs] memory = torch.zeros([1, 1, self.policy.m_size]) next_obs = [obs.unsqueeze(0) for obs in next_obs] value_estimates, next_memory = self.policy.actor_critic.critic_pass( current_obs, memory, sequence_length=batch.num_experiences) next_value_estimate, _ = self.policy.actor_critic.critic_pass( next_obs, next_memory, sequence_length=1) for name, estimate in value_estimates.items(): value_estimates[name] = ModelUtils.to_numpy(estimate) next_value_estimate[name] = ModelUtils.to_numpy( next_value_estimate[name]) if done: for k in next_value_estimate: if not self.reward_signals[k].ignore_done: next_value_estimate[k] = 0.0 return value_estimates, next_value_estimate
def evaluate(self, decision_requests: DecisionSteps, global_agent_ids: List[str]) -> Dict[str, Any]: """ Evaluates policy for the agent experiences provided. :param global_agent_ids: :param decision_requests: DecisionStep object containing inputs. :return: Outputs from network as defined by self.inference_dict. """ vec_vis_obs, masks = self._split_decision_step(decision_requests) vec_obs = [torch.as_tensor(vec_vis_obs.vector_observations)] vis_obs = [ torch.as_tensor(vis_ob) for vis_ob in vec_vis_obs.visual_observations ] memories = torch.as_tensor( self.retrieve_memories(global_agent_ids)).unsqueeze(0) run_out = {} with torch.no_grad(): action, log_probs, entropy, memories = self.sample_actions( vec_obs, vis_obs, masks=masks, memories=memories) run_out["action"] = ModelUtils.to_numpy(action) run_out["pre_action"] = ModelUtils.to_numpy(action) # Todo - make pre_action difference run_out["log_probs"] = ModelUtils.to_numpy(log_probs) run_out["entropy"] = ModelUtils.to_numpy(entropy) run_out["learning_rate"] = 0.0 if self.use_recurrent: run_out["memory_out"] = ModelUtils.to_numpy(memories).squeeze(0) return run_out
def to_log_probs_tuple(self) -> LogProbsTuple: """ Returns a LogProbsTuple. Only adds if tensor is not None. Otherwise, LogProbsTuple uses a default. """ log_probs_tuple = LogProbsTuple() if self.continuous_tensor is not None: continuous = ModelUtils.to_numpy(self.continuous_tensor) log_probs_tuple.add_continuous(continuous) if self.discrete_list is not None: discrete = ModelUtils.to_numpy(self.discrete_tensor) log_probs_tuple.add_discrete(discrete) return log_probs_tuple
def to_action_tuple(self, clip: bool = False) -> ActionTuple: """ Returns an ActionTuple """ action_tuple = ActionTuple() if self.continuous_tensor is not None: _continuous_tensor = self.continuous_tensor if clip: _continuous_tensor = torch.clamp(_continuous_tensor, -3, 3) / 3 continuous = ModelUtils.to_numpy(_continuous_tensor) action_tuple.add_continuous(continuous) if self.discrete_list is not None: discrete = ModelUtils.to_numpy(self.discrete_tensor[:, 0, :]) action_tuple.add_discrete(discrete) return action_tuple
def evaluate(self, mini_batch: AgentBuffer) -> np.ndarray: with torch.no_grad(): estimates, _ = self._discriminator_network.compute_estimate( mini_batch, use_vail_noise=False) return ModelUtils.to_numpy( -torch.log(1.0 - estimates.squeeze(dim=1) * (1.0 - self._discriminator_network.EPSILON)))
def get_trajectory_value_estimates( self, batch: AgentBuffer, next_obs: List[np.ndarray], done: bool) -> Tuple[Dict[str, np.ndarray], Dict[str, float]]: vector_obs = [ModelUtils.list_to_tensor(batch["vector_obs"])] if self.policy.use_vis_obs: visual_obs = [] for idx, _ in enumerate( self.policy.actor_critic.network_body.visual_processors): visual_ob = ModelUtils.list_to_tensor(batch["visual_obs%d" % idx]) visual_obs.append(visual_ob) else: visual_obs = [] memory = torch.zeros([1, 1, self.policy.m_size]) vec_vis_obs = SplitObservations.from_observations(next_obs) next_vec_obs = [ ModelUtils.list_to_tensor( vec_vis_obs.vector_observations).unsqueeze(0) ] next_vis_obs = [ ModelUtils.list_to_tensor(_vis_ob).unsqueeze(0) for _vis_ob in vec_vis_obs.visual_observations ] value_estimates, next_memory = self.policy.actor_critic.critic_pass( vector_obs, visual_obs, memory, sequence_length=batch.num_experiences) next_value_estimate, _ = self.policy.actor_critic.critic_pass( next_vec_obs, next_vis_obs, next_memory, sequence_length=1) for name, estimate in value_estimates.items(): value_estimates[name] = ModelUtils.to_numpy(estimate) next_value_estimate[name] = ModelUtils.to_numpy( next_value_estimate[name]) if done: for k in next_value_estimate: if not self.reward_signals[k].ignore_done: next_value_estimate[k] = 0.0 return value_estimates, next_value_estimate
def test_next_state_prediction(behavior_spec: BehaviorSpec, seed: int) -> None: np.random.seed(seed) torch.manual_seed(seed) curiosity_settings = CuriositySettings(32, 0.1) curiosity_rp = CuriosityRewardProvider(behavior_spec, curiosity_settings) buffer = create_agent_buffer(behavior_spec, 5) for _ in range(100): curiosity_rp.update(buffer) prediction = curiosity_rp._network.predict_next_state(buffer)[0] target = curiosity_rp._network.get_next_state(buffer)[0] error = float(ModelUtils.to_numpy(torch.mean((prediction - target) ** 2))) assert error < 0.001
def evaluate(self, mini_batch: AgentBuffer) -> np.ndarray: with torch.no_grad(): rewards = ModelUtils.to_numpy( self._network.compute_reward(mini_batch)) rewards = np.minimum(rewards, 1.0 / self.strength) return rewards * self._has_updated_once
def _evaluate_by_sequence( self, tensor_obs: List[torch.Tensor], initial_memory: np.ndarray ) -> Tuple[Dict[str, torch.Tensor], AgentBufferField, torch.Tensor]: """ Evaluate a trajectory sequence-by-sequence, assembling the result. This enables us to get the intermediate memories for the critic. :param tensor_obs: A List of tensors of shape (trajectory_len, <obs_dim>) that are the agent's observations for this trajectory. :param initial_memory: The memory that preceeds this trajectory. Of shape (1,1,<mem_size>), i.e. what is returned as the output of a MemoryModules. :return: A Tuple of the value estimates as a Dict of [name, tensor], an AgentBufferField of the initial memories to be used during value function update, and the final memory at the end of the trajectory. """ num_experiences = tensor_obs[0].shape[0] all_next_memories = AgentBufferField() # In the buffer, the 1st sequence are the ones that are padded. So if seq_len = 3 and # trajectory is of length 10, the 1st sequence is [pad,pad,obs]. # Compute the number of elements in this padded seq. leftover = num_experiences % self.policy.sequence_length # Compute values for the potentially truncated initial sequence seq_obs = [] first_seq_len = self.policy.sequence_length for _obs in tensor_obs: if leftover > 0: first_seq_len = leftover first_seq_obs = _obs[0:first_seq_len] seq_obs.append(first_seq_obs) # For the first sequence, the initial memory should be the one at the # beginning of this trajectory. for _ in range(first_seq_len): all_next_memories.append(ModelUtils.to_numpy(initial_memory.squeeze())) init_values, _mem = self.critic.critic_pass( seq_obs, initial_memory, sequence_length=first_seq_len ) all_values = { signal_name: [init_values[signal_name]] for signal_name in init_values.keys() } # Evaluate other trajectories, carrying over _mem after each # trajectory for seq_num in range( 1, math.ceil((num_experiences) / (self.policy.sequence_length)) ): seq_obs = [] for _ in range(self.policy.sequence_length): all_next_memories.append(ModelUtils.to_numpy(_mem.squeeze())) for _obs in tensor_obs: start = seq_num * self.policy.sequence_length - ( self.policy.sequence_length - leftover ) end = (seq_num + 1) * self.policy.sequence_length - ( self.policy.sequence_length - leftover ) seq_obs.append(_obs[start:end]) values, _mem = self.critic.critic_pass( seq_obs, _mem, sequence_length=self.policy.sequence_length ) for signal_name, _val in values.items(): all_values[signal_name].append(_val) # Create one tensor per reward signal all_value_tensors = { signal_name: torch.cat(value_list, dim=0) for signal_name, value_list in all_values.items() } next_mem = _mem return all_value_tensors, all_next_memories, next_mem
def get_trajectory_value_estimates( self, batch: AgentBuffer, next_obs: List[np.ndarray], done: bool, agent_id: str = "", ) -> Tuple[Dict[str, np.ndarray], Dict[str, float], Optional[AgentBufferField]]: """ Get value estimates and memories for a trajectory, in batch form. :param batch: An AgentBuffer that consists of a trajectory. :param next_obs: the next observation (after the trajectory). Used for boostrapping if this is not a termiinal trajectory. :param done: Set true if this is a terminal trajectory. :param agent_id: Agent ID of the agent that this trajectory belongs to. :returns: A Tuple of the Value Estimates as a Dict of [name, np.ndarray(trajectory_len)], the final value estimate as a Dict of [name, float], and optionally (if using memories) an AgentBufferField of initial critic memories to be used during update. """ n_obs = len(self.policy.behavior_spec.observation_specs) if agent_id in self.critic_memory_dict: memory = self.critic_memory_dict[agent_id] else: memory = ( torch.zeros((1, 1, self.critic.memory_size)) if self.policy.use_recurrent else None ) # Convert to tensors current_obs = [ ModelUtils.list_to_tensor(obs) for obs in ObsUtil.from_buffer(batch, n_obs) ] next_obs = [ModelUtils.list_to_tensor(obs) for obs in next_obs] next_obs = [obs.unsqueeze(0) for obs in next_obs] # If we're using LSTM, we want to get all the intermediate memories. all_next_memories: Optional[AgentBufferField] = None # To prevent memory leak and improve performance, evaluate with no_grad. with torch.no_grad(): if self.policy.use_recurrent: ( value_estimates, all_next_memories, next_memory, ) = self._evaluate_by_sequence(current_obs, memory) else: value_estimates, next_memory = self.critic.critic_pass( current_obs, memory, sequence_length=batch.num_experiences ) # Store the memory for the next trajectory. This should NOT have a gradient. self.critic_memory_dict[agent_id] = next_memory next_value_estimate, _ = self.critic.critic_pass( next_obs, next_memory, sequence_length=1 ) for name, estimate in value_estimates.items(): value_estimates[name] = ModelUtils.to_numpy(estimate) next_value_estimate[name] = ModelUtils.to_numpy(next_value_estimate[name]) if done: for k in next_value_estimate: if not self.reward_signals[k].ignore_done: next_value_estimate[k] = 0.0 if agent_id in self.critic_memory_dict: self.critic_memory_dict.pop(agent_id) return value_estimates, next_value_estimate, all_next_memories
def get_trajectory_and_baseline_value_estimates( self, batch: AgentBuffer, next_obs: List[np.ndarray], next_groupmate_obs: List[List[np.ndarray]], done: bool, agent_id: str = "", ) -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray], Dict[str, float], Optional[AgentBufferField], Optional[AgentBufferField], ]: """ Get value estimates, baseline estimates, and memories for a trajectory, in batch form. :param batch: An AgentBuffer that consists of a trajectory. :param next_obs: the next observation (after the trajectory). Used for boostrapping if this is not a termiinal trajectory. :param next_groupmate_obs: the next observations from other members of the group. :param done: Set true if this is a terminal trajectory. :param agent_id: Agent ID of the agent that this trajectory belongs to. :returns: A Tuple of the Value Estimates as a Dict of [name, np.ndarray(trajectory_len)], the baseline estimates as a Dict, the final value estimate as a Dict of [name, float], and optionally (if using memories) an AgentBufferField of initial critic and baseline memories to be used during update. """ n_obs = len(self.policy.behavior_spec.observation_specs) current_obs = ObsUtil.from_buffer(batch, n_obs) groupmate_obs = GroupObsUtil.from_buffer(batch, n_obs) current_obs = [ModelUtils.list_to_tensor(obs) for obs in current_obs] groupmate_obs = [[ ModelUtils.list_to_tensor(obs) for obs in _groupmate_obs ] for _groupmate_obs in groupmate_obs] groupmate_actions = AgentAction.group_from_buffer(batch) next_obs = [ModelUtils.list_to_tensor(obs) for obs in next_obs] next_obs = [obs.unsqueeze(0) for obs in next_obs] next_groupmate_obs = [ ModelUtils.list_to_tensor_list(_list_obs) for _list_obs in next_groupmate_obs ] # Expand dimensions of next critic obs next_groupmate_obs = [[_obs.unsqueeze(0) for _obs in _list_obs] for _list_obs in next_groupmate_obs] if agent_id in self.value_memory_dict: # The agent_id should always be in both since they are added together _init_value_mem = self.value_memory_dict[agent_id] _init_baseline_mem = self.baseline_memory_dict[agent_id] else: _init_value_mem = (torch.zeros((1, 1, self.critic.memory_size)) if self.policy.use_recurrent else None) _init_baseline_mem = (torch.zeros((1, 1, self.critic.memory_size)) if self.policy.use_recurrent else None) all_obs = ([current_obs] + groupmate_obs if groupmate_obs is not None else [current_obs]) all_next_value_mem: Optional[AgentBufferField] = None all_next_baseline_mem: Optional[AgentBufferField] = None with torch.no_grad(): if self.policy.use_recurrent: ( value_estimates, baseline_estimates, all_next_value_mem, all_next_baseline_mem, next_value_mem, next_baseline_mem, ) = self._evaluate_by_sequence_team( current_obs, groupmate_obs, groupmate_actions, _init_value_mem, _init_baseline_mem, ) else: value_estimates, next_value_mem = self.critic.critic_pass( all_obs, _init_value_mem, sequence_length=batch.num_experiences) groupmate_obs_and_actions = (groupmate_obs, groupmate_actions) baseline_estimates, next_baseline_mem = self.critic.baseline( current_obs, groupmate_obs_and_actions, _init_baseline_mem, sequence_length=batch.num_experiences, ) # Store the memory for the next trajectory self.value_memory_dict[agent_id] = next_value_mem self.baseline_memory_dict[agent_id] = next_baseline_mem all_next_obs = ([next_obs] + next_groupmate_obs if next_groupmate_obs is not None else [next_obs]) next_value_estimates, _ = self.critic.critic_pass(all_next_obs, next_value_mem, sequence_length=1) for name, estimate in baseline_estimates.items(): baseline_estimates[name] = ModelUtils.to_numpy(estimate) for name, estimate in value_estimates.items(): value_estimates[name] = ModelUtils.to_numpy(estimate) # the base line and V shpuld not be on the same done flag for name, estimate in next_value_estimates.items(): next_value_estimates[name] = ModelUtils.to_numpy(estimate) if done: for k in next_value_estimates: if not self.reward_signals[k].ignore_done: next_value_estimates[k][-1] = 0.0 return ( value_estimates, baseline_estimates, next_value_estimates, all_next_value_mem, all_next_baseline_mem, )
def _evaluate_by_sequence_team( self, self_obs: List[torch.Tensor], obs: List[List[torch.Tensor]], actions: List[AgentAction], init_value_mem: torch.Tensor, init_baseline_mem: torch.Tensor, ) -> Tuple[Dict[str, torch.Tensor], Dict[ str, torch.Tensor], AgentBufferField, AgentBufferField, torch.Tensor, torch.Tensor, ]: """ Evaluate a trajectory sequence-by-sequence, assembling the result. This enables us to get the intermediate memories for the critic. :param tensor_obs: A List of tensors of shape (trajectory_len, <obs_dim>) that are the agent's observations for this trajectory. :param initial_memory: The memory that preceeds this trajectory. Of shape (1,1,<mem_size>), i.e. what is returned as the output of a MemoryModules. :return: A Tuple of the value estimates as a Dict of [name, tensor], an AgentBufferField of the initial memories to be used during value function update, and the final memory at the end of the trajectory. """ num_experiences = self_obs[0].shape[0] all_next_value_mem = AgentBufferField() all_next_baseline_mem = AgentBufferField() # In the buffer, the 1st sequence are the ones that are padded. So if seq_len = 3 and # trajectory is of length 10, the 1st sequence is [pad,pad,obs]. # Compute the number of elements in this padded seq. leftover = num_experiences % self.policy.sequence_length # Compute values for the potentially truncated initial sequence first_seq_len = leftover if leftover > 0 else self.policy.sequence_length self_seq_obs = [] groupmate_seq_obs = [] groupmate_seq_act = [] seq_obs = [] for _self_obs in self_obs: first_seq_obs = _self_obs[0:first_seq_len] seq_obs.append(first_seq_obs) self_seq_obs.append(seq_obs) for groupmate_obs, groupmate_action in zip(obs, actions): seq_obs = [] for _obs in groupmate_obs: first_seq_obs = _obs[0:first_seq_len] seq_obs.append(first_seq_obs) groupmate_seq_obs.append(seq_obs) _act = groupmate_action.slice(0, first_seq_len) groupmate_seq_act.append(_act) # For the first sequence, the initial memory should be the one at the # beginning of this trajectory. for _ in range(first_seq_len): all_next_value_mem.append( ModelUtils.to_numpy(init_value_mem.squeeze())) all_next_baseline_mem.append( ModelUtils.to_numpy(init_baseline_mem.squeeze())) all_seq_obs = self_seq_obs + groupmate_seq_obs init_values, _value_mem = self.critic.critic_pass( all_seq_obs, init_value_mem, sequence_length=first_seq_len) all_values = { signal_name: [init_values[signal_name]] for signal_name in init_values.keys() } groupmate_obs_and_actions = (groupmate_seq_obs, groupmate_seq_act) init_baseline, _baseline_mem = self.critic.baseline( self_seq_obs[0], groupmate_obs_and_actions, init_baseline_mem, sequence_length=first_seq_len, ) all_baseline = { signal_name: [init_baseline[signal_name]] for signal_name in init_baseline.keys() } # Evaluate other trajectories, carrying over _mem after each # trajectory for seq_num in range( 1, math.ceil( (num_experiences) / (self.policy.sequence_length))): for _ in range(self.policy.sequence_length): all_next_value_mem.append( ModelUtils.to_numpy(_value_mem.squeeze())) all_next_baseline_mem.append( ModelUtils.to_numpy(_baseline_mem.squeeze())) start = seq_num * self.policy.sequence_length - ( self.policy.sequence_length - leftover) end = (seq_num + 1) * self.policy.sequence_length - ( self.policy.sequence_length - leftover) self_seq_obs = [] groupmate_seq_obs = [] groupmate_seq_act = [] seq_obs = [] for _self_obs in self_obs: seq_obs.append(_obs[start:end]) self_seq_obs.append(seq_obs) for groupmate_obs, team_action in zip(obs, actions): seq_obs = [] for (_obs, ) in groupmate_obs: first_seq_obs = _obs[start:end] seq_obs.append(first_seq_obs) groupmate_seq_obs.append(seq_obs) _act = team_action.slice(start, end) groupmate_seq_act.append(_act) all_seq_obs = self_seq_obs + groupmate_seq_obs values, _value_mem = self.critic.critic_pass( all_seq_obs, _value_mem, sequence_length=self.policy.sequence_length) all_values = { signal_name: [init_values[signal_name]] for signal_name in values.keys() } groupmate_obs_and_actions = (groupmate_seq_obs, groupmate_seq_act) baselines, _baseline_mem = self.critic.baseline( self_seq_obs[0], groupmate_obs_and_actions, _baseline_mem, sequence_length=first_seq_len, ) all_baseline = { signal_name: [baselines[signal_name]] for signal_name in baselines.keys() } # Create one tensor per reward signal all_value_tensors = { signal_name: torch.cat(value_list, dim=0) for signal_name, value_list in all_values.items() } all_baseline_tensors = { signal_name: torch.cat(baseline_list, dim=0) for signal_name, baseline_list in all_baseline.items() } next_value_mem = _value_mem next_baseline_mem = _baseline_mem return ( all_value_tensors, all_baseline_tensors, all_next_value_mem, all_next_baseline_mem, next_value_mem, next_baseline_mem, )
def _evaluate_by_sequence_team( self, self_obs: List[torch.Tensor], obs: List[List[torch.Tensor]], actions: List[AgentAction], init_value_mem: torch.Tensor, init_baseline_mem: torch.Tensor, ) -> Tuple[Dict[str, torch.Tensor], Dict[ str, torch.Tensor], AgentBufferField, AgentBufferField, torch.Tensor, torch.Tensor, ]: """ Evaluate a trajectory sequence-by-sequence, assembling the result. This enables us to get the intermediate memories for the critic. :param tensor_obs: A List of tensors of shape (trajectory_len, <obs_dim>) that are the agent's observations for this trajectory. :param initial_memory: The memory that preceeds this trajectory. Of shape (1,1,<mem_size>), i.e. what is returned as the output of a MemoryModules. :return: A Tuple of the value estimates as a Dict of [name, tensor], an AgentBufferField of the initial memories to be used during value function update, and the final memory at the end of the trajectory. """ num_experiences = self_obs[0].shape[0] all_next_value_mem = AgentBufferField() all_next_baseline_mem = AgentBufferField() # When using LSTM, we need to divide the trajectory into sequences of equal length. Sometimes, # that division isn't even, and we must pad the leftover sequence. # In the buffer, the last sequence are the ones that are padded. So if seq_len = 3 and # trajectory is of length 10, the last sequence is [obs,pad,pad]. # Compute the number of elements in this padded seq. leftover_seq_len = num_experiences % self.policy.sequence_length all_values: Dict[str, List[np.ndarray]] = defaultdict(list) all_baseline: Dict[str, List[np.ndarray]] = defaultdict(list) _baseline_mem = init_baseline_mem _value_mem = init_value_mem # Evaluate other trajectories, carrying over _mem after each # trajectory for seq_num in range(num_experiences // self.policy.sequence_length): for _ in range(self.policy.sequence_length): all_next_value_mem.append( ModelUtils.to_numpy(_value_mem.squeeze())) all_next_baseline_mem.append( ModelUtils.to_numpy(_baseline_mem.squeeze())) start = seq_num * self.policy.sequence_length end = (seq_num + 1) * self.policy.sequence_length self_seq_obs = [] groupmate_seq_obs = [] groupmate_seq_act = [] seq_obs = [] for _self_obs in self_obs: seq_obs.append(_self_obs[start:end]) self_seq_obs.append(seq_obs) for groupmate_obs, groupmate_action in zip(obs, actions): seq_obs = [] for _obs in groupmate_obs: sliced_seq_obs = _obs[start:end] seq_obs.append(sliced_seq_obs) groupmate_seq_obs.append(seq_obs) _act = groupmate_action.slice(start, end) groupmate_seq_act.append(_act) all_seq_obs = self_seq_obs + groupmate_seq_obs values, _value_mem = self.critic.critic_pass( all_seq_obs, _value_mem, sequence_length=self.policy.sequence_length) for signal_name, _val in values.items(): all_values[signal_name].append(_val) groupmate_obs_and_actions = (groupmate_seq_obs, groupmate_seq_act) baselines, _baseline_mem = self.critic.baseline( self_seq_obs[0], groupmate_obs_and_actions, _baseline_mem, sequence_length=self.policy.sequence_length, ) for signal_name, _val in baselines.items(): all_baseline[signal_name].append(_val) # Compute values for the potentially truncated initial sequence if leftover_seq_len > 0: self_seq_obs = [] groupmate_seq_obs = [] groupmate_seq_act = [] seq_obs = [] for _self_obs in self_obs: last_seq_obs = _self_obs[-leftover_seq_len:] seq_obs.append(last_seq_obs) self_seq_obs.append(seq_obs) for groupmate_obs, groupmate_action in zip(obs, actions): seq_obs = [] for _obs in groupmate_obs: last_seq_obs = _obs[-leftover_seq_len:] seq_obs.append(last_seq_obs) groupmate_seq_obs.append(seq_obs) _act = groupmate_action.slice( len(_obs) - leftover_seq_len, len(_obs)) groupmate_seq_act.append(_act) # For the last sequence, the initial memory should be the one at the # beginning of this trajectory. seq_obs = [] for _ in range(leftover_seq_len): all_next_value_mem.append( ModelUtils.to_numpy(_value_mem.squeeze())) all_next_baseline_mem.append( ModelUtils.to_numpy(_baseline_mem.squeeze())) all_seq_obs = self_seq_obs + groupmate_seq_obs last_values, _value_mem = self.critic.critic_pass( all_seq_obs, _value_mem, sequence_length=leftover_seq_len) for signal_name, _val in last_values.items(): all_values[signal_name].append(_val) groupmate_obs_and_actions = (groupmate_seq_obs, groupmate_seq_act) last_baseline, _baseline_mem = self.critic.baseline( self_seq_obs[0], groupmate_obs_and_actions, _baseline_mem, sequence_length=leftover_seq_len, ) for signal_name, _val in last_baseline.items(): all_baseline[signal_name].append(_val) # Create one tensor per reward signal all_value_tensors = { signal_name: torch.cat(value_list, dim=0) for signal_name, value_list in all_values.items() } all_baseline_tensors = { signal_name: torch.cat(baseline_list, dim=0) for signal_name, baseline_list in all_baseline.items() } next_value_mem = _value_mem next_baseline_mem = _baseline_mem return ( all_value_tensors, all_baseline_tensors, all_next_value_mem, all_next_baseline_mem, next_value_mem, next_baseline_mem, )
def _evaluate_by_sequence( self, tensor_obs: List[torch.Tensor], initial_memory: torch.Tensor ) -> Tuple[Dict[str, torch.Tensor], AgentBufferField, torch.Tensor]: """ Evaluate a trajectory sequence-by-sequence, assembling the result. This enables us to get the intermediate memories for the critic. :param tensor_obs: A List of tensors of shape (trajectory_len, <obs_dim>) that are the agent's observations for this trajectory. :param initial_memory: The memory that preceeds this trajectory. Of shape (1,1,<mem_size>), i.e. what is returned as the output of a MemoryModules. :return: A Tuple of the value estimates as a Dict of [name, tensor], an AgentBufferField of the initial memories to be used during value function update, and the final memory at the end of the trajectory. """ num_experiences = tensor_obs[0].shape[0] all_next_memories = AgentBufferField() # When using LSTM, we need to divide the trajectory into sequences of equal length. Sometimes, # that division isn't even, and we must pad the leftover sequence. # When it is added to the buffer, the last sequence will be padded. So if seq_len = 3 and # trajectory is of length 10, the last sequence is [obs,pad,pad] once it is added to the buffer. # Compute the number of elements in this sequence that will end up being padded. leftover_seq_len = num_experiences % self.policy.sequence_length all_values: Dict[str, List[np.ndarray]] = defaultdict(list) _mem = initial_memory # Evaluate other trajectories, carrying over _mem after each # trajectory for seq_num in range(num_experiences // self.policy.sequence_length): seq_obs = [] for _ in range(self.policy.sequence_length): all_next_memories.append(ModelUtils.to_numpy(_mem.squeeze())) start = seq_num * self.policy.sequence_length end = (seq_num + 1) * self.policy.sequence_length for _obs in tensor_obs: seq_obs.append(_obs[start:end]) values, _mem = self.critic.critic_pass( seq_obs, _mem, sequence_length=self.policy.sequence_length) for signal_name, _val in values.items(): all_values[signal_name].append(_val) # Compute values for the potentially truncated last sequence. Note that this # sequence isn't padded yet, but will be. seq_obs = [] if leftover_seq_len > 0: for _obs in tensor_obs: last_seq_obs = _obs[-leftover_seq_len:] seq_obs.append(last_seq_obs) # For the last sequence, the initial memory should be the one at the # end of this trajectory. for _ in range(leftover_seq_len): all_next_memories.append(ModelUtils.to_numpy(_mem.squeeze())) last_values, _mem = self.critic.critic_pass( seq_obs, _mem, sequence_length=leftover_seq_len) for signal_name, _val in last_values.items(): all_values[signal_name].append(_val) # Create one tensor per reward signal all_value_tensors = { signal_name: torch.cat(value_list, dim=0) for signal_name, value_list in all_values.items() } next_mem = _mem return all_value_tensors, all_next_memories, next_mem