Add BwdModelView instead of shallow copy for bwd model

Author: etaoxing

rewarped/autograd.py

@@ -6,16 +6,18 @@
 
 
 # for checkpointing method
-def assign_tensors(x, x_out, names, tensors):
+def assign_tensors(x, x_out, names, tensors, view=False):
     # need to assign b/c state_0, state_1 cannot be swapped
+    # if view=True, then x == x_out except for tensors given by names, so we can skip assigning some
     # TODO: Add fn to get wp.array attributes instead of vars(..)
-    for name in vars(x):
-        if name in names:
-            continue
-        attr = getattr(x, name)
-        if isinstance(attr, wp.array):
-            wp_array = getattr(x_out, name)
-            wp_array.assign(attr)
+    if not view:
+        for name in vars(x):
+            if name in names:
+                continue
+            attr = getattr(x, name)
+            if isinstance(attr, wp.array):
+                wp_array = getattr(x_out, name)
+                wp_array.assign(attr)
     for name, tensor in zip(names, tensors, strict=True):
         # assert not torch.isnan(tensor).any(), print("NaN tensor", name)
         wp_array = getattr(x_out, name)
@@ -115,7 +117,7 @@ def forward(
                     finally:
                         tape.bwd_update_graph = wp.capture_end()
 
-            assign_tensors(model, model_bwd, model_tensors_names, model_tensors)
+            assign_tensors(model, model_bwd, model_tensors_names, model_tensors, view=True)
             assign_tensors(state_in, state_in_bwd, state_tensors_names, state_tensors)
             assign_tensors(control, control_bwd, control_tensors_names, control_tensors)
             wp.capture_launch(tape.update_graph)
@@ -197,7 +199,7 @@ def backward(ctx, *adj_tensors):
 
         if use_graph_capture:
             # checkpointing method
-            assign_tensors(model, model_bwd, model_tensors_names, model_tensors)
+            assign_tensors(model, model_bwd, model_tensors_names, model_tensors, view=True)
             assign_tensors(state_in, state_in_bwd, state_tensors_names, state_tensors)
             assign_tensors(control, control_bwd, control_tensors_names, control_tensors)
             wp.capture_launch(tape.update_graph)

rewarped/warp_env.py

@@ -1,5 +1,3 @@
-import copy
-
 import numpy as np
 import torch
 from gym import spaces
@@ -52,6 +50,23 @@ def __getattr__(self, name):
         return wp.to_torch(getattr(self.data, name))
 
 
+class BwdModelView:
+    def __init__(self, model, model_tensors_names):
+        self.model = model
+        self.model_tensors_names = model_tensors_names
+
+        self.bwd_tensors = {}
+        for k in model_tensors_names:
+            v = getattr(model, k)
+            v = wp.zeros_like(v, requires_grad=v.requires_grad)
+            self.bwd_tensors[k] = v
+
+    def __getattr__(self, name):
+        if name in self.model_tensors_names:
+            return self.bwd_tensors[name]
+        return getattr(self.model, name)
+
+
 class WarpEnv(Environment):
     r"""Base class for gym-like Warp environments that builds on `Environment`.
 
@@ -290,13 +305,8 @@ def init_sim(self):
             if self.use_graph_capture:
                 self.tape = wp.Tape()  # persistent tape for graph capture
 
-                # shallow copy
-                # TODO: need a better day to have separate copies when not using graph capture (for randomization)
-                self.model_bwd = copy.copy(self.model)
-                for k in self.model_tensors_names:
-                    v = getattr(self.model_bwd, k)
-                    v = wp.zeros_like(v, requires_grad=self.requires_grad)
-                    setattr(self.model_bwd, k, v)
+                # shallow copy of model with new arrays for `model_tensors`
+                self.model_bwd = BwdModelView(self.model, self.model_tensors_names)
 
                 self.state_0_bwd = self.model.state(copy="zeros")
                 self.state_1_bwd = self.model.state(copy="zeros")