feature(rjy): add HAPPO algorithm

ding/model/common/head.py

@@ -1103,7 +1103,7 @@ class ReparameterizationHead(nn.Module):
         ``__init__``, ``forward``.
     """
 
-    default_sigma_type = ['fixed', 'independent', 'conditioned']
+    default_sigma_type = ['fixed', 'independent', 'conditioned', 'happo']
     default_bound_type = ['tanh', None]
 
     def __init__(
@@ -1155,6 +1155,11 @@ def __init__(
             self.log_sigma_param = nn.Parameter(torch.zeros(1, output_size))
         elif self.sigma_type == 'conditioned':
             self.log_sigma_layer = nn.Linear(hidden_size, output_size)
+        elif self.sigma_type == 'happo':
+            self.sigma_x_coef = 1.
+            self.sigma_y_coef = 0.5
+            # This parameter (x_coef, y_coef) refers to the HAPPO paper
+            self.log_sigma_param = nn.Parameter(torch.ones(1, output_size) * self.sigma_x_coef)
 
     def forward(self, x: torch.Tensor) -> Dict:
         """
@@ -1190,6 +1195,9 @@ def forward(self, x: torch.Tensor) -> Dict:
         elif self.sigma_type == 'conditioned':
             log_sigma = self.log_sigma_layer(x)
             sigma = torch.exp(torch.clamp(log_sigma, -20, 2))
+        elif self.sigma_type == 'happo':
+            log_sigma = self.log_sigma_param + torch.zeros_like(mu)
+            sigma = torch.sigmoid(log_sigma / self.sigma_x_coef) * self.sigma_y_coef
         return {'mu': mu, 'sigma': sigma}
 
 

ding/model/template/__init__.py

@@ -27,3 +27,4 @@
 from .bcq import BCQ
 from .edac import EDAC
 from .ebm import EBM, AutoregressiveEBM
+from .havac import HAVAC

ding/model/template/tests/test_havac.py

@@ -0,0 +1,102 @@
+import pytest
+import torch
+import random
+from ding.torch_utils import is_differentiable
+from ding.model.template import HAVAC
+
+
+@pytest.mark.unittest
+class TestHAVAC:
+
+    def test_havac_rnn_actor(self):
+        # discrete+rnn
+        bs, T = 3, 8
+        obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9
+        agent_num = 5
+        data = {
+            'obs': {
+                'agent_state': torch.randn(T, bs, obs_dim),
+                'global_state': torch.randn(T, bs, global_obs_dim),
+                'action_mask': torch.randint(0, 2, size=(T, bs, action_dim))
+            },
+            'actor_prev_state': [None for _ in range(bs)],
+        }
+        model = HAVAC(
+            agent_obs_shape=obs_dim,
+            global_obs_shape=global_obs_dim,
+            action_shape=action_dim,
+            agent_num=agent_num,
+            use_lstm=True,
+        )
+        agent_idx = random.randint(0, agent_num - 1)
+        output = model(agent_idx, data, mode='compute_actor')
+        assert set(output.keys()) == set(['logit', 'actor_next_state', 'actor_hidden_state'])
+        assert output['logit'].shape == (T, bs, action_dim)
+        assert len(output['actor_next_state']) == bs
+        print(output['actor_next_state'][0]['h'].shape)
+        loss = output['logit'].sum()
+        is_differentiable(loss, model.agent_models[agent_idx].actor)
+
+    def test_havac_rnn_critic(self):
+        # discrete+rnn
+        bs, T = 3, 8
+        obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9
+        agent_num = 5
+        data = {
+            'obs': {
+                'agent_state': torch.randn(T, bs, obs_dim),
+                'global_state': torch.randn(T, bs, global_obs_dim),
+                'action_mask': torch.randint(0, 2, size=(T, bs, action_dim))
+            },
+            'critic_prev_state': [None for _ in range(bs)],
+        }
+        model = HAVAC(
+            agent_obs_shape=obs_dim,
+            global_obs_shape=global_obs_dim,
+            action_shape=action_dim,
+            agent_num=agent_num,
+            use_lstm=True,
+        )
+        agent_idx = random.randint(0, agent_num - 1)
+        output = model(agent_idx, data, mode='compute_critic')
+        assert set(output.keys()) == set(['value', 'critic_next_state', 'critic_hidden_state'])
+        assert output['value'].shape == (T, bs)
+        assert len(output['critic_next_state']) == bs
+        print(output['critic_next_state'][0]['h'].shape)
+        loss = output['value'].sum()
+        is_differentiable(loss, model.agent_models[agent_idx].critic)
+
+    def test_havac_rnn_actor_critic(self):
+        # discrete+rnn
+        bs, T = 3, 8
+        obs_dim, global_obs_dim, action_dim = 32, 32 * 4, 9
+        agent_num = 5
+        data = {
+            'obs': {
+                'agent_state': torch.randn(T, bs, obs_dim),
+                'global_state': torch.randn(T, bs, global_obs_dim),
+                'action_mask': torch.randint(0, 2, size=(T, bs, action_dim))
+            },
+            'actor_prev_state': [None for _ in range(bs)],
+            'critic_prev_state': [None for _ in range(bs)],
+        }
+        model = HAVAC(
+            agent_obs_shape=obs_dim,
+            global_obs_shape=global_obs_dim,
+            action_shape=action_dim,
+            agent_num=agent_num,
+            use_lstm=True,
+        )
+        agent_idx = random.randint(0, agent_num - 1)
+        output = model(agent_idx, data, mode='compute_actor_critic')
+        assert set(output.keys()) == set(['logit', 'actor_next_state', 'actor_hidden_state',
+                                          'value', 'critic_next_state', 'critic_hidden_state'])
+        assert output['logit'].shape == (T, bs, action_dim)
+        assert output['value'].shape == (T, bs)
+        loss = output['logit'].sum() + output['value'].sum()
+        is_differentiable(loss, model.agent_models[agent_idx])
+
+
+# test_havac_rnn_actor()
+# test_havac_rnn_critic()
+# test_havac_rnn_actor_critic()

ding/policy/__init__.py

@@ -55,3 +55,4 @@
 # new-type policy
 from .ppof import PPOFPolicy
 from .prompt_pg import PromptPGPolicy
+from .happo import HAPPOPolicy

ding/policy/command_mode_policy_instance.py

@@ -51,6 +51,7 @@
 from .edac import EDACPolicy
 from .prompt_pg import PromptPGPolicy
 from .plan_diffuser import PDPolicy
+from .happo import HAPPOPolicy
 
 
 class EpsCommandModePolicy(CommandModePolicy):
@@ -186,6 +187,11 @@ class PPOCommandModePolicy(PPOPolicy, DummyCommandModePolicy):
     pass
 
 
+@POLICY_REGISTRY.register('happo_command')
+class HAPPOCommandModePolicy(HAPPOPolicy, DummyCommandModePolicy):
+    pass
+
+
 @POLICY_REGISTRY.register('ppo_stdim_command')
 class PPOSTDIMCommandModePolicy(PPOSTDIMPolicy, DummyCommandModePolicy):
     pass