ignoring precommit, grabbing baseline/critic mems from buffer in trainer

Author: andrewcoh

ml-agents/mlagents/trainers/buffer.py

@@ -35,6 +35,7 @@ class BufferKey(enum.Enum):
  MASKS = "masks"
  MEMORY = "memory"
  CRITIC_MEMORY = "critic_memory"
+ BASELINE_MEMORY = "coma_baseline_memory"
  PREV_ACTION = "prev_action"
 
  ADVANTAGES = "advantages"

ml-agents/mlagents/trainers/coma/optimizer_torch.py

@@ -264,6 +264,23 @@ def update(self, batch: AgentBuffer, num_sequences: int) -> Dict[str, float]:
  ]
  if len(memories) > 0:
  memories = torch.stack(memories).unsqueeze(0)
+ value_memories = [
+ ModelUtils.list_to_tensor(batch[BufferKey.CRITIC_MEMORY][i])
+ for i in range(
+ 0, len(batch[BufferKey.CRITIC_MEMORY]), self.policy.sequence_length
+ )
+ ]
+
+ baseline_memories = [
+ ModelUtils.list_to_tensor(batch[BufferKey.BASELINE_MEMORY][i])
+ for i in range(
+ 0, len(batch[BufferKey.BASELINE_MEMORY]), self.policy.sequence_length
+ )
+ ]
+
+ if len(value_memories) > 0:
+ value_memories = torch.stack(value_memories).unsqueeze(0)
+ baseline_memories = torch.stack(baseline_memories).unsqueeze(0)
 
  log_probs, entropy = self.policy.evaluate_actions(
  current_obs,
@@ -274,13 +291,15 @@ def update(self, batch: AgentBuffer, num_sequences: int) -> Dict[str, float]:
  )
  all_obs = [current_obs] + group_obs
  values, _ = self.critic.critic_pass(
- all_obs, memories=memories, sequence_length=self.policy.sequence_length
+ all_obs,
+ memories=value_memories,
+ sequence_length=self.policy.sequence_length,
  )
  baselines, _ = self.critic.baseline(
  [current_obs],
  group_obs,
  group_actions,
- memories=memories,
+ memories=baseline_memories,
  sequence_length=self.policy.sequence_length,
  )
  old_log_probs = ActionLogProbs.from_buffer(batch).flatten()
@@ -380,7 +399,7 @@ def _evaluate_by_sequence_team(
 
  for team_obs, team_action in zip(obs, actions):
  seq_obs = []
- for (_obs,) in team_obs:
+ for _obs in team_obs:
  first_seq_obs = _obs[0:first_seq_len]
  seq_obs.append(first_seq_obs)
  team_seq_obs.append(seq_obs)
@@ -534,7 +553,12 @@ def get_trajectory_and_baseline_value_estimates(
  _init_value_mem = self.value_memory_dict[agent_id]
  _init_baseline_mem = self.baseline_memory_dict[agent_id]
  else:
- memory = (
+ _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
@@ -544,19 +568,19 @@ def get_trajectory_and_baseline_value_estimates(
  all_next_value_mem: Optional[AgentBufferField] = None
  all_next_baseline_mem: Optional[AgentBufferField] = None
  if self.policy.use_recurrent:
- value_estimates, baseline_estimates, all_next_value_mem, all_next_baseline_mem, next_value_mem, next_baseline_mem = self.critic._evaluate_by_sequence_team(
+ 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, team_obs, team_actions, _init_value_mem, _init_baseline_mem
  )
  else:
  value_estimates, value_mem = self.critic.critic_pass(
- all_obs, memory, sequence_length=batch.num_experiences
+ all_obs, _init_value_mem, sequence_length=batch.num_experiences
  )
 
  baseline_estimates, baseline_mem = self.critic.baseline(
  [current_obs],
  team_obs,
  team_actions,
- memory,
+ _init_baseline_mem,
  sequence_length=batch.num_experiences,
  )
  # Store the memory for the next trajectory

ml-agents/mlagents/trainers/coma/trainer.py

@@ -80,6 +80,8 @@ def _process_trajectory(self, trajectory: Trajectory) -> None:
  value_estimates,
  baseline_estimates,
  value_next,
+ value_memories,
+ baseline_memories,
  ) = self.optimizer.get_trajectory_and_baseline_value_estimates(
  agent_buffer_trajectory,
  trajectory.next_obs,
@@ -89,6 +91,10 @@ def _process_trajectory(self, trajectory: Trajectory) -> None:
  and not trajectory.interrupted,
  )
 
+ if value_memories is not None:
+ agent_buffer_trajectory[BufferKey.CRITIC_MEMORY].set(value_memories)
+ agent_buffer_trajectory[BufferKey.BASELINE_MEMORY].set(baseline_memories)
+
  for name, v in value_estimates.items():
  agent_buffer_trajectory[RewardSignalUtil.value_estimates_key(name)].extend(
  v

ml-agents/mlagents/trainers/optimizer/torch_optimizer.py

@@ -56,7 +56,7 @@ def create_reward_signals(self, reward_signal_configs):
  )
 
  def _evaluate_by_sequence(
- self, tensor_obs: List[torch.Tensor], initial_memory: np.ndarray
+ 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