[coma2] Make group extrinsic reward part of extrinsic (#5033)

* Make group extrinsic part of extrinsic * Fix test and init * Fix tests and bug * Add baseline loss to TensorBoard

Author: Ervin T

config/ppo/PushBlockCollab.yaml

@@ -16,11 +16,11 @@ behaviors:
  num_layers: 2
  vis_encode_type: simple
  reward_signals:
- group:
+ extrinsic:
  gamma: 0.99
  strength: 1.0
  keep_checkpoints: 5
- max_steps: 20000000 #2000000
+ max_steps: 20000000
  time_horizon: 64
  summary_freq: 60000
  threaded: true

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

@@ -14,7 +14,13 @@
 from mlagents_envs.base_env import ObservationSpec, ActionSpec
 from mlagents.trainers.policy.torch_policy import TorchPolicy
 from mlagents.trainers.optimizer.torch_optimizer import TorchOptimizer
-from mlagents.trainers.settings import TrainerSettings, PPOSettings
+from mlagents.trainers.settings import (
+ ExtrinsicSettings,
+ RewardSignalSettings,
+ RewardSignalType,
+ TrainerSettings,
+ PPOSettings,
+)
 from mlagents.trainers.torch.networks import Critic, MultiInputNetworkBody
 from mlagents.trainers.torch.decoders import ValueHeads
 from mlagents.trainers.torch.agent_action import AgentAction
@@ -23,6 +29,10 @@
 from mlagents.trainers.trajectory import ObsUtil, GroupObsUtil
 from mlagents.trainers.settings import NetworkSettings
 
+from mlagents_envs.logging_util import get_logger
+
+logger = get_logger(__name__)
+
 
 class TorchCOMAOptimizer(TorchOptimizer):
  class COMAValueNetwork(torch.nn.Module, Critic):
@@ -157,6 +167,24 @@ def __init__(self, policy: TorchPolicy, trainer_settings: TrainerSettings):
  self.value_memory_dict: Dict[str, torch.Tensor] = {}
  self.baseline_memory_dict: Dict[str, torch.Tensor] = {}
 
+ def create_reward_signals(
+ self, reward_signal_configs: Dict[RewardSignalType, RewardSignalSettings]
+ ) -> None:
+ """
+ Create reward signals. Override default to provide warnings for Curiosity and
+ GAIL, and make sure Extrinsic adds team rewards.
+ :param reward_signal_configs: Reward signal config.
+ """
+ for reward_signal, settings in reward_signal_configs.items():
+ if reward_signal != RewardSignalType.EXTRINSIC:
+ logger.warning(
+ f"Reward Signal {reward_signal.value} is not supported with the COMA2 trainer; \
+ results may be unexpected."
+ )
+ elif isinstance(settings, ExtrinsicSettings):
+ settings.add_groupmate_rewards = True
+ super().create_reward_signals(reward_signal_configs)
+
  @property
  def critic(self):
  return self._critic
@@ -335,6 +363,7 @@ def update(self, batch: AgentBuffer, num_sequences: int) -> Dict[str, float]:
  # TODO: After PyTorch is default, change to something more correct.
  "Losses/Policy Loss": torch.abs(policy_loss).item(),
  "Losses/Value Loss": value_loss.item(),
+ "Losses/Baseline Loss": baseline_loss.item(),
  "Policy/Learning Rate": decay_lr,
  "Policy/Epsilon": decay_eps,
  "Policy/Beta": decay_bet,

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

@@ -10,7 +10,11 @@
 
 from mlagents.trainers.policy.torch_policy import TorchPolicy
 from mlagents.trainers.optimizer import Optimizer
-from mlagents.trainers.settings import TrainerSettings
+from mlagents.trainers.settings import (
+ TrainerSettings,
+ RewardSignalSettings,
+ RewardSignalType,
+)
 from mlagents.trainers.torch.utils import ModelUtils
 
 
@@ -44,7 +48,9 @@ def critic(self):
  def update(self, batch: AgentBuffer, num_sequences: int) -> Dict[str, float]:
  pass
 
- def create_reward_signals(self, reward_signal_configs):
+ def create_reward_signals(
+ self, reward_signal_configs: Dict[RewardSignalType, RewardSignalSettings]
+ ) -> None:
  """
  Create reward signals
  :param reward_signal_configs: Reward signal config.

ml-agents/mlagents/trainers/settings.py

@@ -165,15 +165,13 @@ def _reward_signal_steps_per_update_default(self):
 # INTRINSIC REWARD SIGNALS #############################################################
 class RewardSignalType(Enum):
  EXTRINSIC: str = "extrinsic"
- GROUP_EXTRINSIC: str = "group"
  GAIL: str = "gail"
  CURIOSITY: str = "curiosity"
  RND: str = "rnd"
 
  def to_settings(self) -> type:
  _mapping = {
- RewardSignalType.EXTRINSIC: RewardSignalSettings,
- RewardSignalType.GROUP_EXTRINSIC: RewardSignalSettings,
+ RewardSignalType.EXTRINSIC: ExtrinsicSettings,
  RewardSignalType.GAIL: GAILSettings,
  RewardSignalType.CURIOSITY: CuriositySettings,
  RewardSignalType.RND: RNDSettings,
@@ -217,6 +215,12 @@ def structure(d: Mapping, t: type) -> Any:
  return d_final
 
 
+@attr.s(auto_attribs=True)
+class ExtrinsicSettings(RewardSignalSettings):
+ # For use with COMA2. Add groupmate rewards to the final extrinsic reward.
+ add_groupmate_rewards = False
+
+
 @attr.s(auto_attribs=True)
 class GAILSettings(RewardSignalSettings):
  learning_rate: float = 3e-4
@@ -625,7 +629,7 @@ def _set_default_hyperparameters(self):
 
  network_settings: NetworkSettings = attr.ib(factory=NetworkSettings)
  reward_signals: Dict[RewardSignalType, RewardSignalSettings] = attr.ib(
- factory=lambda: {RewardSignalType.EXTRINSIC: RewardSignalSettings()}
+ factory=lambda: {RewardSignalType.EXTRINSIC: ExtrinsicSettings()}
  )
  init_path: Optional[str] = None
  keep_checkpoints: int = 5

ml-agents/mlagents/trainers/tests/torch/test_coma.py

@@ -4,7 +4,7 @@
 import attr
 
 from mlagents.trainers.coma.optimizer_torch import TorchCOMAOptimizer
-from mlagents.trainers.settings import RewardSignalSettings, RewardSignalType
+from mlagents.trainers.settings import ExtrinsicSettings, RewardSignalType
 
 from mlagents.trainers.policy.torch_policy import TorchPolicy
 from mlagents.trainers.tests import mock_brain as mb
@@ -49,7 +49,9 @@ def create_test_coma_optimizer(dummy_config, use_rnn, use_discrete, use_visual):
 
  trainer_settings = attr.evolve(dummy_config)
  trainer_settings.reward_signals = {
- RewardSignalType.GROUP_EXTRINSIC: RewardSignalSettings(strength=1.0, gamma=0.99)
+ RewardSignalType.EXTRINSIC: ExtrinsicSettings(
+ strength=1.0, gamma=0.99, add_groupmate_rewards=True
+ )
  }
 
  trainer_settings.network_settings.memory = (
@@ -122,7 +124,11 @@ def test_coma_get_value_estimates(dummy_config, rnn, visual, discrete):
  max_step_complete=True,
  num_other_agents_in_group=NUM_AGENTS,
  )
- value_estimates, baseline_estimates, next_value_estimates = optimizer.get_trajectory_and_baseline_value_estimates(
+ (
+ value_estimates,
+ baseline_estimates,
+ next_value_estimates,
+ ) = optimizer.get_trajectory_and_baseline_value_estimates(
  trajectory.to_agentbuffer(),
  trajectory.next_obs,
  trajectory.next_group_obs,
@@ -138,7 +144,11 @@ def test_coma_get_value_estimates(dummy_config, rnn, visual, discrete):
  # if all_memories is not None:
  # assert len(all_memories) == 15
 
- value_estimates, baseline_estimates, next_value_estimates = optimizer.get_trajectory_and_baseline_value_estimates(
+ (
+ value_estimates,
+ baseline_estimates,
+ next_value_estimates,
+ ) = optimizer.get_trajectory_and_baseline_value_estimates(
  trajectory.to_agentbuffer(),
  trajectory.next_obs,
  trajectory.next_group_obs,
@@ -150,7 +160,11 @@ def test_coma_get_value_estimates(dummy_config, rnn, visual, discrete):
 
  # Check if we ignore terminal states properly
  optimizer.reward_signals["group"].use_terminal_states = False
- value_estimates, baseline_estimates, next_value_estimates = optimizer.get_trajectory_and_baseline_value_estimates(
+ (
+ value_estimates,
+ baseline_estimates,
+ next_value_estimates,
+ ) = optimizer.get_trajectory_and_baseline_value_estimates(
  trajectory.to_agentbuffer(),
  trajectory.next_obs,
  trajectory.next_group_obs,

ml-agents/mlagents/trainers/tests/torch/test_reward_providers/test_extrinsic.py

@@ -1,10 +1,12 @@
+from mlagents.trainers.buffer import BufferKey
 import pytest
+import numpy as np
 from mlagents.trainers.torch.components.reward_providers import (
  ExtrinsicRewardProvider,
  create_reward_provider,
 )
 from mlagents_envs.base_env import BehaviorSpec, ActionSpec
-from mlagents.trainers.settings import RewardSignalSettings, RewardSignalType
+from mlagents.trainers.settings import ExtrinsicSettings, RewardSignalType
 from mlagents.trainers.tests.torch.test_reward_providers.utils import (
  create_agent_buffer,
 )
@@ -27,7 +29,7 @@
  ],
 )
 def test_construction(behavior_spec: BehaviorSpec) -> None:
- settings = RewardSignalSettings()
+ settings = ExtrinsicSettings()
  settings.gamma = 0.2
  extrinsic_rp = ExtrinsicRewardProvider(behavior_spec, settings)
  assert extrinsic_rp.gamma == 0.2
@@ -46,7 +48,7 @@ def test_construction(behavior_spec: BehaviorSpec) -> None:
  ],
 )
 def test_factory(behavior_spec: BehaviorSpec) -> None:
- settings = RewardSignalSettings()
+ settings = ExtrinsicSettings()
  extrinsic_rp = create_reward_provider(
  RewardSignalType.EXTRINSIC, behavior_spec, settings
  )
@@ -67,7 +69,24 @@ def test_factory(behavior_spec: BehaviorSpec) -> None:
 )
 def test_reward(behavior_spec: BehaviorSpec, reward: float) -> None:
  buffer = create_agent_buffer(behavior_spec, 1000, reward)
- settings = RewardSignalSettings()
+ settings = ExtrinsicSettings()
  extrinsic_rp = ExtrinsicRewardProvider(behavior_spec, settings)
  generated_rewards = extrinsic_rp.evaluate(buffer)
  assert (generated_rewards == reward).all()
+
+ # Test group rewards. Rewards should be double of the environment rewards, but shouldn't count
+ # the groupmate rewards.
+ buffer[BufferKey.GROUP_REWARD] = buffer[BufferKey.ENVIRONMENT_REWARDS]
+ # 2 agents with identical rewards
+ buffer[BufferKey.GROUPMATE_REWARDS].set(
+ [np.ones(1, dtype=np.float32) * reward] * 2
+ for _ in range(buffer.num_experiences)
+ )
+ generated_rewards = extrinsic_rp.evaluate(buffer)
+ assert (generated_rewards == 2 * reward).all()
+
+ # Test groupmate rewards. Total reward should be indiv_reward + 2 * teammate_reward + group_reward
+ settings.add_groupmate_rewards = True
+ extrinsic_rp = ExtrinsicRewardProvider(behavior_spec, settings)
+ generated_rewards = extrinsic_rp.evaluate(buffer)
+ assert (generated_rewards == 4 * reward).all()

ml-agents/mlagents/trainers/torch/components/reward_providers/__init__.py

@@ -4,9 +4,6 @@
 from mlagents.trainers.torch.components.reward_providers.extrinsic_reward_provider import ( # noqa F401
  ExtrinsicRewardProvider,
 )
-from mlagents.trainers.torch.components.reward_providers.group_extrinsic_reward_provider import ( # noqa F401
- GroupExtrinsicRewardProvider,
-)
 from mlagents.trainers.torch.components.reward_providers.curiosity_reward_provider import ( # noqa F401
  CuriosityRewardProvider,
 )

ml-agents/mlagents/trainers/torch/components/reward_providers/extrinsic_reward_provider.py

@@ -5,11 +5,42 @@
 from mlagents.trainers.torch.components.reward_providers.base_reward_provider import (
  BaseRewardProvider,
 )
+from mlagents_envs.base_env import BehaviorSpec
+from mlagents.trainers.settings import ExtrinsicSettings
 
 
 class ExtrinsicRewardProvider(BaseRewardProvider):
+ """
+ Evaluates extrinsic reward. For single-agent, this equals the individual reward
+ given to the agent. For the COMA2 algorithm, we want not only the individual reward
+ but also the team and the individual rewards of the other agents.
+ """
+
+ def __init__(self, specs: BehaviorSpec, settings: ExtrinsicSettings) -> None:
+ super().__init__(specs, settings)
+ self._add_groupmate_rewards = settings.add_groupmate_rewards
+
  def evaluate(self, mini_batch: AgentBuffer) -> np.ndarray:
- return np.array(mini_batch[BufferKey.ENVIRONMENT_REWARDS], dtype=np.float32)
+ indiv_rewards = np.array(
+ mini_batch[BufferKey.ENVIRONMENT_REWARDS], dtype=np.float32
+ )
+ total_rewards = indiv_rewards
+ if (
+ BufferKey.GROUPMATE_REWARDS in mini_batch
+ and BufferKey.GROUP_REWARD in mini_batch
+ ):
+ if self._add_groupmate_rewards:
+ groupmate_rewards_list = mini_batch[BufferKey.GROUPMATE_REWARDS]
+ groupmate_rewards_sum = np.array(
+ [sum(_rew) for _rew in groupmate_rewards_list], dtype=np.float32
+ )
+ total_rewards += groupmate_rewards_sum
+ group_rewards = np.array(
+ mini_batch[BufferKey.GROUP_REWARD], dtype=np.float32
+ )
+ # Add all the group rewards to the individual rewards
+ total_rewards += group_rewards
+ return total_rewards
 
  def update(self, mini_batch: AgentBuffer) -> Dict[str, np.ndarray]:
  return {}

ml-agents/mlagents/trainers/torch/components/reward_providers/group_extrinsic_reward_provider.py

@@ -15,9 +15,6 @@
 from mlagents.trainers.torch.components.reward_providers.gail_reward_provider import (
  GAILRewardProvider,
 )
-from mlagents.trainers.torch.components.reward_providers.group_extrinsic_reward_provider import (
- GroupExtrinsicRewardProvider,
-)
 from mlagents.trainers.torch.components.reward_providers.rnd_reward_provider import (
  RNDRewardProvider,
 )
@@ -26,7 +23,6 @@
 
 NAME_TO_CLASS: Dict[RewardSignalType, Type[BaseRewardProvider]] = {
  RewardSignalType.EXTRINSIC: ExtrinsicRewardProvider,
- RewardSignalType.GROUP_EXTRINSIC: GroupExtrinsicRewardProvider,
  RewardSignalType.CURIOSITY: CuriosityRewardProvider,
  RewardSignalType.GAIL: GAILRewardProvider,
  RewardSignalType.RND: RNDRewardProvider,