forked from NM512/dreamerv3-torch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
exploration.py
119 lines (109 loc) · 4.27 KB
/
exploration.py
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import torch
from torch import nn
from torch import distributions as torchd
import models
import networks
import tools
class Random(nn.Module):
def __init__(self, config):
self._config = config
def actor(self, feat):
shape = feat.shape[:-1] + [self._config.num_actions]
if self._config.actor_dist == "onehot":
return tools.OneHotDist(torch.zeros(shape))
else:
ones = torch.ones(shape)
return tools.ContDist(torchd.uniform.Uniform(-ones, ones))
def train(self, start, context):
return None, {}
# class Plan2Explore(tools.Module):
class Plan2Explore(nn.Module):
def __init__(self, config, world_model, reward=None):
super(Plan2Explore, self).__init__()
self._config = config
self._use_amp = True if config.precision == 16 else False
self._reward = reward
self._behavior = models.ImagBehavior(config, world_model)
self.actor = self._behavior.actor
if config.dyn_discrete:
feat_size = config.dyn_stoch * config.dyn_discrete + config.dyn_deter
stoch = config.dyn_stoch * config.dyn_discrete
else:
feat_size = config.dyn_stoch + config.dyn_deter
stoch = config.dyn_stoch
size = {
"embed": 32 * config.cnn_depth,
"stoch": stoch,
"deter": config.dyn_deter,
"feat": config.dyn_stoch + config.dyn_deter,
}[self._config.disag_target]
kw = dict(
inp_dim=feat_size, # pytorch version
shape=size,
layers=config.disag_layers,
units=config.disag_units,
act=config.act,
)
self._networks = nn.ModuleList(
[networks.DenseHead(**kw) for _ in range(config.disag_models)]
)
kw = dict(wd=config.weight_decay, opt=config.opt, use_amp=self._use_amp)
self._model_opt = tools.Optimizer(
"explorer",
self.parameters(),
config.model_lr,
config.opt_eps,
config.grad_clip,
**kw
)
def train(self, start, context, data):
with tools.RequiresGrad(self):
metrics = {}
stoch = start["stoch"]
if self._config.dyn_discrete:
stoch = torch.reshape(
stoch, (stoch.shape[:-2] + ((stoch.shape[-2] * stoch.shape[-1]),))
)
target = {
"embed": context["embed"],
"stoch": stoch,
"deter": start["deter"],
"feat": context["feat"],
}[self._config.disag_target]
inputs = context["feat"]
if self._config.disag_action_cond:
inputs = torch.concat([inputs, data["action"]], -1)
metrics.update(self._train_ensemble(inputs, target))
metrics.update(self._behavior._train(start, self._intrinsic_reward)[-1])
return None, metrics
def _intrinsic_reward(self, feat, state, action):
inputs = feat
if self._config.disag_action_cond:
inputs = torch.concat([inputs, action], -1)
preds = torch.cat(
[head(inputs, torch.float32).mode()[None] for head in self._networks], 0
)
disag = torch.mean(torch.std(preds, 0), -1)[..., None]
if self._config.disag_log:
disag = torch.log(disag)
reward = self._config.expl_intr_scale * disag
if self._config.expl_extr_scale:
reward += torch.cast(
self._config.expl_extr_scale * self._reward(feat, state, action),
torch.float32,
)
return reward
def _train_ensemble(self, inputs, targets):
with torch.cuda.amp.autocast(self._use_amp):
if self._config.disag_offset:
targets = targets[:, self._config.disag_offset :]
inputs = inputs[:, : -self._config.disag_offset]
targets = targets.detach()
inputs = inputs.detach()
preds = [head(inputs) for head in self._networks]
likes = torch.cat(
[torch.mean(pred.log_prob(targets))[None] for pred in preds], 0
)
loss = -torch.mean(likes)
metrics = self._model_opt(loss, self.parameters())
return metrics