Update sb3_imitation.py

Refactor for loop
edbeeching · Jan 29, 2024 · 02a87b7 · 02a87b7
1 parent 187106c
commit 02a87b7
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/examples/sb3_imitation.py b/examples/sb3_imitation.py
@@ -115,11 +115,11 @@ def close_env():
     with open(file_path, "r") as file:
         data = json.load(file)
 
-    for i in range(0, len(data)):
+    for traj in data:
         trajectories.append(
             imitation.data.rollout.types.Trajectory(
-                obs=np.array(data[i][0]),
-                acts=np.array(data[i][1]),
+                obs=np.array(traj[0]),
+                acts=np.array(traj[1]),
                 infos=None,
                 terminal=True,
             )