minor cleanup

VT-Collab · Sep 6, 2024 · 9064572 · 9064572
1 parent 5faa523
commit 9064572
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/generate_ablation_demos.py b/generate_ablation_demos.py
@@ -114,7 +114,7 @@ def parse_args():
                 grip = target[-1]
                 action = np.hstack((linear, angular, [grip]))
 
-                robot_state, _,_, info = sim.step(action)
+                robot_state, _, _,_, info = sim.step(action)
                 state = np.hstack((robot_state[:3], robot_state[-1]))
                 object_state = np.array(info['object_positions'][object_of_interest][:3]).squeeze() # remove orientation of object
 

diff --git a/helper_functions.py b/helper_functions.py
@@ -370,7 +370,7 @@ def simulate(args, traj, object_of_interest, verbose=False):
 
         robot_state, _,_, info = sim.step(action)
         while not grip == robot_state[-1]: # wait for gripper to open/close
-            robot_state, _,_, info = sim.step(action)
+            robot_state, _,_, _, info = sim.step(action)
         state = np.hstack((robot_state[:3], robot_state[-1]))
         object_state = np.array(info['object_positions'][object_of_interest][:3]) # remove orientation of object