您最多选择25个主题
主题必须以中文或者字母或数字开头，可以包含连字符 (-)，并且长度不得超过35个字符
22 行
784 B
22 行
784 B
from typing import List, Dict
|
|
|
|
from mlagents.torch_utils import torch, nn
|
|
from mlagents.trainers.torch.layers import linear_layer
|
|
|
|
|
|
class ValueHeads(nn.Module):
|
|
def __init__(self, stream_names: List[str], input_size: int, output_size: int = 1):
|
|
super().__init__()
|
|
self.stream_names = stream_names
|
|
_value_heads = {}
|
|
|
|
for name in stream_names:
|
|
value = linear_layer(input_size, output_size)
|
|
_value_heads[name] = value
|
|
self.value_heads = nn.ModuleDict(_value_heads)
|
|
|
|
def forward(self, hidden: torch.Tensor) -> Dict[str, torch.Tensor]:
|
|
value_outputs = {}
|
|
for stream_name, head in self.value_heads.items():
|
|
value_outputs[stream_name] = head(hidden).squeeze(-1)
|
|
return value_outputs
|