-
Notifications
You must be signed in to change notification settings - Fork 5
Expand file tree
/
Copy pathppg.py
More file actions
executable file
·62 lines (54 loc) · 2.17 KB
/
Copy pathppg.py
File metadata and controls
executable file
·62 lines (54 loc) · 2.17 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
from typing import Optional, Dict, Union
import copy
import torch
import torch.nn as nn
from ding.utils import SequenceType, MODEL_REGISTRY
from .vac import VAC
@MODEL_REGISTRY.register('ppg')
class PPG(nn.Module):
mode = ['compute_actor', 'compute_critic', 'compute_actor_critic']
def __init__(
self,
obs_shape: Union[int, SequenceType],
action_shape: Union[int, SequenceType],
share_encoder: bool = True,
continuous: bool = False,
encoder_hidden_size_list: SequenceType = [128, 128, 64],
actor_head_hidden_size: int = 64,
actor_head_layer_num: int = 2,
critic_head_hidden_size: int = 64,
critic_head_layer_num: int = 1,
activation: Optional[nn.Module] = nn.ReLU(),
norm_type: Optional[str] = None,
) -> None:
super(PPG, self).__init__()
self.actor_critic = VAC(
obs_shape, action_shape, share_encoder, continuous, encoder_hidden_size_list, actor_head_hidden_size,
actor_head_layer_num, critic_head_hidden_size, critic_head_layer_num, activation, norm_type
)
self.aux_critic = copy.deepcopy(self.actor_critic.critic)
def forward(self, inputs: Union[torch.Tensor, Dict], mode: str) -> Dict:
assert mode in self.mode, "not support forward mode: {}/{}".format(mode, self.mode)
return getattr(self, mode)(inputs)
def compute_actor(self, x: torch.Tensor) -> Dict:
"""
ReturnsKeys:
- necessary: ``logit``
"""
return self.actor_critic(x, mode='compute_actor')
def compute_critic(self, x: torch.Tensor) -> Dict:
"""
ReturnsKeys:
- necessary: ``value``
"""
x = self.aux_critic[0](x) # encoder
x = self.aux_critic[1](x) # head
return {'value': x['pred']}
def compute_actor_critic(self, x: torch.Tensor) -> Dict:
"""
.. note::
``compute_actor_critic`` interface aims to save computation when shares encoder
ReturnsKeys:
- necessary: ``value``, ``logit``
"""
return self.actor_critic(x, mode='compute_actor_critic')