HBBC部署到代码中

Env/hbbc_actor_critic.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+
+
+def _get_activation(name: str):
+    name = (name or "elu").lower()
+    mapping = {
+        "elu": nn.ELU,
+        "relu": nn.ReLU,
+        "tanh": nn.Tanh,
+        "leakyrelu": nn.LeakyReLU,
+    }
+    if name not in mapping:
+        raise ValueError(f"Unsupported activation: {name}")
+    return mapping[name]()
+
+
+class ActorCritic(nn.Module):
+    """Minimal HBBC ActorCritic for inference-only deployment."""
+
+    def __init__(
+        self,
+        num_actor_obs=18,
+        num_critic_obs=18,
+        num_actions=2,
+        latent_c_dim=4,
+        latent_eps_dim=6,
+        use_style_latent=True,
+        actor_hidden_dims=None,
+        activation="elu",
+    ):
+        super().__init__()
+        _ = num_critic_obs  # kept for checkpoint compatibility
+        if actor_hidden_dims is None:
+            actor_hidden_dims = [512, 256, 128]
+
+        act_fn = _get_activation(activation)
+        self.latent_c_dim = int(latent_c_dim)
+        self.latent_eps_dim = int(latent_eps_dim)
+        self.use_style_latent = bool(use_style_latent)
+
+        layers = [nn.Linear(num_actor_obs, actor_hidden_dims[0]), act_fn]
+        for i in range(len(actor_hidden_dims) - 1):
+            layers.append(nn.Linear(actor_hidden_dims[i], actor_hidden_dims[i + 1]))
+            layers.append(_get_activation(activation))
+        self.actor_trunk = nn.Sequential(*layers)
+        self.actor_head = nn.Linear(actor_hidden_dims[-1], num_actions)
+
+        if self.use_style_latent:
+            self.style_trunk = nn.Sequential(
+                nn.Linear(self.latent_eps_dim, 512),
+                _get_activation(activation),
+                nn.Linear(512, 256),
+                _get_activation(activation),
+                nn.Linear(256, 128),
+                _get_activation(activation),
+            )
+            self.style_head = nn.Linear(128, self.latent_eps_dim)
+            self.style_activation = torch.tanh
+
+    def act_inference(self, observations: torch.Tensor) -> torch.Tensor:
+        if self.use_style_latent:
+            obs = observations[..., :-(self.latent_c_dim + self.latent_eps_dim)]
+            eps = observations[..., -self.latent_c_dim - self.latent_eps_dim:-self.latent_c_dim]
+            c = observations[..., -self.latent_c_dim:]
+            eps = self.style_activation(self.style_head(self.style_trunk(eps)))
+            observations = torch.cat([obs, eps, c], dim=-1)
+        embedding = self.actor_trunk(observations)
+        return self.actor_head(embedding)