70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
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)
|