move requires_grad of env_params (fix locuslab#5)

only enables grad recording after changing env_params's device
brunompacheco · Jan 13, 2023 · 7585126 · 7585126
1 parent 5e40c02
commit 7585126
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/imitation_nonconvex/il_exp.py b/imitation_nonconvex/il_exp.py
@@ -135,17 +135,16 @@ def __init__(self, **kwargs):
             if self.learn_dx:
                 if self.env_name == 'pendulum':
                     self.env_params = torch.tensor(
-                        (15., 3., 0.5), requires_grad=True)
+                        (15., 3., 0.5))
                 elif self.env_name == 'cartpole':
                     self.env_params = torch.tensor(
-                        (9.8, 3.0, 0.1, 1.0), requires_grad=True)
+                        (9.8, 3.0, 0.1, 1.0))
                 elif self.env_name == 'pendulum-complex':
                     # self.env_params = torch.tensor(
-                    #     (15., 3., 0.5), requires_grad=True)
+                    #     (15., 3., 0.5))
                     torch.manual_seed(self.seed)
                     self.env_params = torch.tensor((5., 1., 1.)) + \
                         torch.tensor((3., 1., 1.))*(torch.rand(3)-0.5)
-                    self.env_params.requires_grad_()
 
                     # n_hidden = 256
                     # self.extra_dx = NNDynamics(
@@ -155,6 +154,7 @@ def __init__(self, **kwargs):
             else:
                 self.env_params = self.env.true_dx.params
             self.env_params = self.env_params.to(self.device)
+            self.env_params.requires_grad_()
         else:
             assert False