hongshaorou/tdmpc2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor pi outputs

Browse Source
This commit is contained in:
Nicklas Hansen
2024-12-25 12:02:33 -08:00
parent 804f9b3949
commit db1865334e
1 changed files with 1 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
tdmpc2/common/world_model.py
View File

@@ -154,13 +154,12 @@ class WorldModel(nn.Module):
action = mean + eps * log_std.exp() action = mean + eps * log_std.exp()
mean, action, log_prob = math.squash(mean, action, log_prob) mean, action, log_prob = math.squash(mean, action, log_prob)
entropy_scale = scaled_log_prob / (log_prob + 1e-8)
info = TensorDict({ info = TensorDict({
"mean": mean, "mean": mean,
"log_std": log_std, "log_std": log_std,
"action_prob": 1., "action_prob": 1.,
"entropy": -log_prob, "entropy": -log_prob,
"scaled_entropy": -log_prob * entropy_scale, "scaled_entropy": -scaled_log_prob,
}) })
return action, info return action, info