wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-26 13:37:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RLlib] SAC: log_alpha not being learnt when on GPU. (#11298)

Browse Source
This commit is contained in:
Sven Mika
2020-10-12 13:48:44 -07:00
committed by GitHub
parent 7dcfd258cd
commit f5e2cda68a
2 changed files with 20 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rllib/agents/sac/sac_torch_model.py
+3 -4
View File
@@ -131,10 +131,9 @@ class SACTorchModel(TorchModelV2, nn.Module):
else:
self.twin_q_net = None
self.log_alpha = torch.tensor(
data=[np.log(initial_alpha)],
dtype=torch.float32,
requires_grad=True)
log_alpha = nn.Parameter(
torch.from_numpy(np.array([np.log(initial_alpha)])).float())
self.register_parameter("log_alpha", log_alpha)
# Auto-calculate the target entropy.
if target_entropy is None or target_entropy == "auto":
rllib/agents/sac/tests/test_sac.py
+17 -8
View File
@@ -109,6 +109,7 @@ class TestSAC(unittest.TestCase):
"_model.0.weight",
"default_policy/value_out/bias": "_value_branch."
"_model.0.bias",
"default_policy/log_alpha": "log_alpha",
# Target net.
"default_policy/sequential_2/action_1/kernel": "action_model."
"action_0._model.0.weight",
@@ -130,6 +131,7 @@ class TestSAC(unittest.TestCase):
"_model.0.weight",
"default_policy/value_out_1/bias": "_value_branch."
"_model.0.bias",
"default_policy/log_alpha_1": "log_alpha",
}
env = SimpleEnv
@@ -186,7 +188,7 @@ class TestSAC(unittest.TestCase):
# Actually convert to torch tensors (by accessing everything).
input_ = policy._lazy_tensor_dict(input_)
input_ = {k: input_[k] for k in input_.keys()}
log_alpha = policy.model.log_alpha.detach().numpy()[0]
log_alpha = policy.model.log_alpha.detach().cpu().numpy()[0]
# Only run the expectation once, should be the same anyways
# for all frameworks.
@@ -259,7 +261,7 @@ class TestSAC(unittest.TestCase):
]
for tf_g, torch_g in zip(tf_a_grads, torch_a_grads):
if tf_g.shape != torch_g.shape:
check(tf_g, np.transpose(torch_g))
check(tf_g, np.transpose(torch_g.detach().cpu()))
else:
check(tf_g, torch_g)
@@ -283,7 +285,7 @@ class TestSAC(unittest.TestCase):
torch_c_grads = [v.grad for v in policy.model.q_variables()]
for tf_g, torch_g in zip(tf_c_grads, torch_c_grads):
if tf_g.shape != torch_g.shape:
check(tf_g, np.transpose(torch_g))
check(tf_g, np.transpose(torch_g.detach().cpu()))
else:
check(tf_g, torch_g)
# Compare (unchanged(!) actor grads) with tf ones.
@@ -292,7 +294,7 @@ class TestSAC(unittest.TestCase):
]
for tf_g, torch_g in zip(tf_a_grads, torch_a_grads):
if tf_g.shape != torch_g.shape:
check(tf_g, np.transpose(torch_g))
check(tf_g, np.transpose(torch_g.detach().cpu()))
else:
check(tf_g, torch_g)
@@ -354,7 +356,10 @@ class TestSAC(unittest.TestCase):
tf_var = tf_weights[tf_key]
torch_var = policy.model.state_dict()[map_[tf_key]]
if tf_var.shape != torch_var.shape:
check(tf_var, np.transpose(torch_var), rtol=0.05)
check(
tf_var,
np.transpose(torch_var.detach().cpu()),
rtol=0.05)
else:
check(tf_var, torch_var, rtol=0.05)
# And alpha.
@@ -366,7 +371,10 @@ class TestSAC(unittest.TestCase):
torch_var = policy.target_model.state_dict()[map_[
tf_key]]
if tf_var.shape != torch_var.shape:
check(tf_var, np.transpose(torch_var), rtol=0.05)
check(
tf_var,
np.transpose(torch_var.detach().cpu()),
rtol=0.05)
else:
check(tf_var, torch_var, rtol=0.05)
@@ -510,9 +518,10 @@ class TestSAC(unittest.TestCase):
def _translate_weights_to_torch(self, weights_dict, map_):
model_dict = {
map_[k]: convert_to_torch_tensor(
np.transpose(v) if re.search("kernel", k) else v)
np.transpose(v) if re.search("kernel", k) else np.array([v])
if re.search("log_alpha", k) else v)
for k, v in weights_dict.items()
if re.search("(sequential(/|_1)|value_out/)", k)
if re.search("(sequential(/|_1)|value_out/|log_alpha)", k)
}
return model_dict