90 lines
3.4 KiB
Python
90 lines
3.4 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from rsl_rl.modules import MLP, EmpiricalNormalization
|
||
|
|
from rsl_rl.utils import resolve_nn_activation
|
||
|
|
|
||
|
|
class Discriminator(nn.Module):
|
||
|
|
"""
|
||
|
|
Simple feedforward neural network as a discriminator for Adversarial Motion Prior.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
input_dim_per_frame: int = 58,
|
||
|
|
input_history_length: int = 1,
|
||
|
|
hidden_dims: list[int] = [256, 128],
|
||
|
|
activation: str = "relu",
|
||
|
|
feature_normalization: bool = False, # if True, normalize input features with EmpiricalNormalization
|
||
|
|
device: str = "cpu"
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
self.device = device
|
||
|
|
curr_in_dim = input_dim_per_frame * input_history_length
|
||
|
|
print("Discriminator input dim:", curr_in_dim)
|
||
|
|
self.frame_size = input_dim_per_frame
|
||
|
|
self.history_length = input_history_length
|
||
|
|
|
||
|
|
self.activation = resolve_nn_activation(activation) #type: ignore
|
||
|
|
|
||
|
|
layers = []
|
||
|
|
for hidden_dim in hidden_dims:
|
||
|
|
layers.append(nn.Linear(curr_in_dim, hidden_dim))
|
||
|
|
layers.append(self.activation) # use resolved activation
|
||
|
|
curr_in_dim = hidden_dim
|
||
|
|
self.model = nn.Sequential(*layers).to(self.device) # type: ignore
|
||
|
|
self.linear_layer = nn.Linear(hidden_dims[-1], 1).to(self.device)
|
||
|
|
|
||
|
|
self.feature_normalization = feature_normalization
|
||
|
|
if self.feature_normalization:
|
||
|
|
self.feature_norm = EmpiricalNormalization(shape=(self.frame_size,)).to(self.device)
|
||
|
|
|
||
|
|
def normalize_input(self, x):
|
||
|
|
if self.feature_normalization:
|
||
|
|
# avoid in-place on a leaf tensor by normalizing frame slices and re-concatenating
|
||
|
|
frames = torch.split(x, self.frame_size, dim=1)
|
||
|
|
norm_frames = [self.feature_norm(frame) for frame in frames]
|
||
|
|
x = torch.cat(norm_frames, dim=1)
|
||
|
|
return x
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
assert self.history_length * self.frame_size == x.shape[1], \
|
||
|
|
f"Input feature dimension {x.shape[1]} does not match expected size {self.history_length * self.frame_size}"
|
||
|
|
if self.feature_normalization:
|
||
|
|
x = self.normalize_input(x)
|
||
|
|
return self.linear_layer(self.model(x)).squeeze(-1)
|
||
|
|
|
||
|
|
# TODO: normalize feature on positive or negative samples?
|
||
|
|
def update_normalization(self, x):
|
||
|
|
if self.feature_normalization:
|
||
|
|
for i in range(self.history_length):
|
||
|
|
start_idx = i * self.frame_size
|
||
|
|
end_idx = (i + 1) * self.frame_size
|
||
|
|
self.feature_norm.update(x[:, start_idx:end_idx])
|
||
|
|
|
||
|
|
@torch.no_grad()
|
||
|
|
def get_amp_reward(self, x):
|
||
|
|
self.eval()
|
||
|
|
d = self.forward(x)
|
||
|
|
r = torch.clamp(1 - 0.25 * (d - 1) ** 2, min=0)
|
||
|
|
self.train()
|
||
|
|
return r
|
||
|
|
|
||
|
|
def compute_grad_pen(self, expert_data, lambda_=10):
|
||
|
|
expert_data.requires_grad = True
|
||
|
|
disc = self.forward(expert_data)
|
||
|
|
ones = torch.ones(disc.size(), device=disc.device)
|
||
|
|
grad = torch.autograd.grad(
|
||
|
|
outputs=disc, inputs=expert_data,
|
||
|
|
grad_outputs=ones, create_graph=True,
|
||
|
|
retain_graph=True, only_inputs=True)[0]
|
||
|
|
|
||
|
|
# Enforce that the grad norm approaches 0.
|
||
|
|
grad_pen = lambda_ * (grad.norm(2, dim=1) - 0).pow(2).mean()
|
||
|
|
return grad_pen
|
||
|
|
|