Updated examples to use optimizer_kwargs dict in forward().

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
     rev: v0.910
     hooks:
       - id: mypy
-        additional_dependencies: [torch==1.9.0, tokenize-rt==3.2.0, types-PyYAML] 
+        additional_dependencies: [torch==1.9.0, tokenize-rt==3.2.0, types-PyYAML, types-mock] 
         args: [--no-strict-optional, --ignore-missing-imports]
         exclude: setup.py
 

examples/backward_modes.py

@@ -79,9 +79,11 @@ def quad_error_fn(optim_vars, aux_vars):
 theseus_optim = th.TheseusLayer(optimizer)
 updated_inputs, info = theseus_optim.forward(
     theseus_inputs,
-    track_best_solution=True,
-    verbose=False,
-    backward_mode=th.BackwardMode.FULL,
+    optimizer_kwargs={
+        "track_best_solution": True,
+        "verbose": False,
+        "backward_mode": th.BackwardMode.FULL,
+    },
 )
 
 # The quadratic \hat y is now fit and we can also use Theseus
@@ -98,9 +100,11 @@ def quad_error_fn(optim_vars, aux_vars):
 # forward again and changing the backward_mode flag.
 updated_inputs, info = theseus_optim.forward(
     theseus_inputs,
-    track_best_solution=True,
-    verbose=False,
-    backward_mode=th.BackwardMode.IMPLICIT,
+    optimizer_kwargs={
+        "track_best_solution": True,
+        "verbose": False,
+        "backward_mode": th.BackwardMode.IMPLICIT,
+    },
 )
 
 da_dx = torch.autograd.grad(updated_inputs["a"], data_x, retain_graph=True)[0].squeeze()
@@ -110,10 +114,12 @@ def quad_error_fn(optim_vars, aux_vars):
 # We can also use truncated unrolling to compute the derivative:
 updated_inputs, info = theseus_optim.forward(
     theseus_inputs,
-    track_best_solution=True,
-    verbose=False,
-    backward_mode=th.BackwardMode.TRUNCATED,
-    backward_num_iterations=5,
+    optimizer_kwargs={
+        "track_best_solution": True,
+        "verbose": False,
+        "backward_mode": th.BackwardMode.TRUNCATED,
+        "backward_num_iterations": 5,
+    },
 )
 
 da_dx = torch.autograd.grad(updated_inputs["a"], data_x, retain_graph=True)[0].squeeze()
@@ -127,8 +133,8 @@ def fit_x(data_x_np):
     theseus_inputs["x"] = (
         torch.from_numpy(data_x_np).float().clone().requires_grad_().unsqueeze(0)
     )
-    updated_inputs, info = theseus_optim.forward(
-        theseus_inputs, track_best_solution=True, verbose=False
+    updated_inputs, _ = theseus_optim.forward(
+        theseus_inputs, optimizer_kwargs={"track_best_solution": True, "verbose": False}
     )
     return updated_inputs["a"].item()
 
@@ -150,9 +156,11 @@ def fit_x(data_x_np):
     start = time.time()
     updated_inputs, info = theseus_optim.forward(
         theseus_inputs,
-        track_best_solution=True,
-        verbose=False,
-        backward_mode=th.BackwardMode.FULL,
+        optimizer_kwargs={
+            "track_best_solution": True,
+            "verbose": False,
+            "backward_mode": th.BackwardMode.FULL,
+        },
     )
     times["fwd"].append(time.time() - start)
 
@@ -164,9 +172,11 @@ def fit_x(data_x_np):
 
     updated_inputs, info = theseus_optim.forward(
         theseus_inputs,
-        track_best_solution=True,
-        verbose=False,
-        backward_mode=th.BackwardMode.IMPLICIT,
+        optimizer_kwargs={
+            "track_best_solution": True,
+            "verbose": False,
+            "backward_mode": th.BackwardMode.IMPLICIT,
+        },
     )
     start = time.time()
     da_dx = torch.autograd.grad(updated_inputs["a"], data_x, retain_graph=True)[
@@ -176,10 +186,12 @@ def fit_x(data_x_np):
 
     updated_inputs, info = theseus_optim.forward(
         theseus_inputs,
-        track_best_solution=True,
-        verbose=False,
-        backward_mode=th.BackwardMode.TRUNCATED,
-        backward_num_iterations=5,
+        optimizer_kwargs={
+            "track_best_solution": True,
+            "verbose": False,
+            "backward_mode": th.BackwardMode.TRUNCATED,
+            "backward_num_iterations": 5,
+        },
     )
     start = time.time()
     da_dx = torch.autograd.grad(updated_inputs["a"], data_x, retain_graph=True)[

examples/motion_planning_2d.py

@@ -149,9 +149,13 @@ def run_learning_loop(cfg):
 
             _, info = motion_planner.layer.forward(
                 planner_inputs,
-                track_best_solution=True,
-                verbose=cfg.verbose,
-                **cfg.optim_params.kwargs,
+                optimizer_kwargs={
+                    **{
+                        "track_best_solution": True,
+                        "verbose": cfg.verbose,
+                    },
+                    **cfg.optim_params.kwargs,
+                },
             )
             if cfg.do_learning and cfg.include_imitation_loss:
                 solution_trajectory = motion_planner.get_trajectory()

examples/state_estimation_2d.py

@@ -293,10 +293,12 @@ def cost_weights_model():
                 print("Initial error:", objective.error_squared_norm().mean().item())
 
         for i in range(inner_loop_iters):
-            theseus_inputs, info = state_estimator.forward(
+            theseus_inputs, _ = state_estimator.forward(
                 theseus_inputs,
-                track_best_solution=True,
-                verbose=epoch % 10 == 0,
+                optimizer_kwargs={
+                    "track_best_solution": True,
+                    "verbose": epoch % 10 == 0,
+                },
             )
             theseus_inputs = run_model(
                 mode_,

examples/tactile_pose_estimation.py

@@ -325,7 +325,9 @@ def run_learning_loop(cfg):
                 (sdf_tensor.data).repeat(batch_size, 1, 1).to(device)
             )
 
-            theseus_inputs, _ = theseus_layer.forward(theseus_inputs, verbose=True)
+            theseus_inputs, _ = theseus_layer.forward(
+                theseus_inputs, optimizer_kwargs={"verbose": True}
+            )
 
             obj_poses_opt = theg.get_tactile_poses_from_values(
                 batch_size=batch_size,

requirements/dev.txt

@@ -5,3 +5,4 @@ nox==2020.8.22
 pre-commit>=2.9.2
 isort>=5.6.4
 types-PyYAML==5.4.3
+types-mock>=4.0.8