[MISC] Enable more example scripts in CI. (#2057)

* Enable more example script in CI.
* Various bug fixes.
* Get around argument pruning bug for ndarray.

Author: duburcqa

examples/coupling/cut_dragon.py

@@ -1,7 +1,5 @@
+import os
 import argparse
-
-import numpy as np
-
 import genesis as gs
 
 
@@ -32,45 +30,45 @@ def main():
             camera_fov=35,
             max_FPS=120,
         ),
-        show_viewer=args.vis,
         vis_options=gs.options.VisOptions(
             visualize_mpm_boundary=True,
             # rendered_envs_idx=[2],
         ),
+        show_viewer=args.vis,
     )
 
     plane = scene.add_entity(
-        material=gs.materials.Rigid(),
         morph=gs.morphs.URDF(file="urdf/plane/plane.urdf", fixed=True),
+        material=gs.materials.Rigid(),
     )
     cutter = scene.add_entity(
         morph=gs.morphs.Mesh(
             file="meshes/cross_cutter.obj",
-            euler=(90, 0, 0),
             scale=0.8,
             pos=(0.0, 0.0, 0.3),
+            euler=(90, 0, 0),
             fixed=True,
             convexify=False,
         ),
         surface=gs.surfaces.Iron(),
     )
     dragon = scene.add_entity(
-        material=gs.materials.MPM.Elastic(sampler="pbs-64"),
         morph=gs.morphs.Mesh(
             file="meshes/dragon/dragon.obj",
             scale=0.007,
             euler=(0, 0, 90),
             pos=(0.3, -0.0, 1.3),
         ),
+        material=gs.materials.MPM.Elastic(sampler="pbs-64"),
         surface=gs.surfaces.Rough(
             color=(0.6, 1.0, 0.8, 1.0),
             vis_mode="particle",
         ),
     )
     scene.build(n_envs=2)
 
-    horizon = 400
-    for i in range(horizon):
+    horizon = 400 if "PYTEST_VERSION" not in os.environ else 5
+    for _ in range(horizon):
         scene.step()
 
 

examples/coupling/fem_cube_linked_with_arm.py

@@ -17,11 +17,9 @@ def main():
         "--substeps", type=int, help="Number of substeps (auto-selected based on solver if not specified)"
     )
     parser.add_argument("--vis", "-v", action="store_true", help="Show visualization GUI")
-    parser.add_argument("-c", "--cpu", action="store_true", default=False)
+    parser.add_argument("-c", "--cpu", action="store_true", default="PYTEST_VERSION" in os.environ)
     args = parser.parse_args()
 
-    args.cpu = args.cpu if "PYTEST_VERSION" not in os.environ else True
-
     gs.init(backend=gs.cpu if args.cpu else gs.gpu, logging_level=None)
 
     if args.solver == "explicit":

examples/coupling/flush_cubes.py

@@ -1,7 +1,6 @@
+import os
 import argparse
-
 import numpy as np
-
 import genesis as gs
 
 
@@ -91,7 +90,7 @@ def main():
     )
     scene.build(n_envs=2)
 
-    horizon = 100
+    horizon = 100 if "PYTEST_VERSION" not in os.environ else 5
     for i in range(horizon):
         emitter1.emit(
             pos=np.array([0.16, -0.4, 0.5]),

examples/differentiable_push.py

@@ -1,4 +1,5 @@
 import argparse
+import os
 
 import torch
 
@@ -89,7 +90,7 @@ def main():
     scene.build(n_envs=2)
 
     ########################## forward + backward twice ##########################
-    horizon = 150
+    horizon = 150 if "PYTEST_VERSION" not in os.environ else 5
     v_list = [gs.tensor([[0.0, 1.0, 0.0], [0.0, 1.0, 0.0]], requires_grad=True) for _ in range(horizon)]
     for _ in range(2):
         scene.reset()

examples/tutorials/IK_motion_planning_grasp.py

@@ -1,5 +1,7 @@
-import genesis as gs
+import os
 import numpy as np
+import genesis as gs
+
 
 ########################## init ##########################
 gs.init(backend=gs.gpu)
@@ -61,7 +63,7 @@
 qpos[-2:] = 0.04
 path = franka.plan_path(
     qpos_goal=qpos,
-    num_waypoints=200,  # 2s duration
+    num_waypoints=200 if "PYTEST_VERSION" not in os.environ else 10,  # 2s duration
 )
 # draw the planned path
 path_debug = scene.draw_debug_path(path, franka)
@@ -75,7 +77,7 @@
 scene.clear_debug_object(path_debug)
 
 # allow robot to reach the last waypoint
-for i in range(100):
+for i in range(100 if "PYTEST_VERSION" not in os.environ else 1):
     scene.step()
 
 # reach
@@ -86,14 +88,14 @@
 )
 print(qpos)
 franka.control_dofs_position(qpos[:-2], motors_dof)
-for i in range(100):
+for i in range(100 if "PYTEST_VERSION" not in os.environ else 1):
     scene.step()
 
 # grasp
 franka.control_dofs_position(qpos[:-2], motors_dof)
 franka.control_dofs_force(np.array([-0.5, -0.5]), fingers_dof)
 
-for i in range(100):
+for i in range(100 if "PYTEST_VERSION" not in os.environ else 1):
     scene.step()
 
 # lift
@@ -104,5 +106,5 @@
 )
 print(qpos)
 franka.control_dofs_position(qpos[:-2], motors_dof)
-for i in range(200):
+for i in range(200 if "PYTEST_VERSION" not in os.environ else 1):
     scene.step()

examples/tutorials/advanced_IK_multilink.py

@@ -1,5 +1,5 @@
+import os
 import numpy as np
-
 import genesis as gs
 
 ########################## init ##########################
@@ -53,7 +53,8 @@
 left_finger = robot.get_link("left_finger")
 right_finger = robot.get_link("right_finger")
 
-for i in range(0, 2000):
+horizon = 2000 if "PYTEST_VERSION" not in os.environ else 5
+for i in range(horizon):
     target_pos_left = center + np.array([np.cos(i / 360 * np.pi), np.sin(i / 360 * np.pi), 0]) * r
     target_pos_right = target_pos_left + np.array([0.0, 0.03, 0])
 

examples/tutorials/advanced_hybrid_robot.py

@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 import genesis as gs
 
@@ -77,15 +79,9 @@
 scene.build()
 
 ########################## run ##########################
+horizon = 1000 if "PYTEST_VERSION" not in os.environ else 5
 scene.reset()
-for i in range(1000):
-    dofs_ctrl = np.array(
-        [
-            1.0 * np.sin(2 * np.pi * i * 0.001),
-        ]
-        * robot.n_dofs
-    )
-
+for i in range(horizon):
+    dofs_ctrl = [1.0 * np.sin(2 * np.pi * i * 0.001)] * robot.n_dofs
     robot.control_dofs_velocity(dofs_ctrl)
-
     scene.step()

examples/tutorials/advanced_muscle.py

@@ -1,3 +1,4 @@
+import os
 import numpy as np
 import genesis as gs
 
@@ -68,10 +69,10 @@
 scene.build(n_envs=2)
 
 ########################## run ##########################
+horizon = 1000 if "PYTEST_VERSION" not in os.environ else 5
 scene.reset()
-for i in range(1000):
-    actu = np.array([0.2 * (0.5 + np.sin(0.01 * np.pi * i))])
-
+for i in range(horizon):
+    actu = [0.2 * (0.5 + np.sin(0.01 * np.pi * i))]
     robot_mpm.set_actuation(actu)
     robot_fem.set_actuation(actu)
     scene.step()

examples/tutorials/advanced_worm.py

@@ -1,5 +1,5 @@
 import math
-
+import os
 import torch
 import genesis as gs
 
@@ -72,7 +72,7 @@
 else:
     raise NotImplementedError
 
-pos_max, pos_min = pos.max(dim=0), pos.min(dim=0)
+pos_max, pos_min = pos.max(dim=0).values, pos.min(dim=0).values
 pos_range = pos_max - pos_min
 
 lu_thr, fh_thr = 0.3, 0.6
@@ -93,8 +93,9 @@
 
 
 ########################## run ##########################
+horizon = 1000 if "PYTEST_VERSION" not in os.environ else 5
 scene.reset()
-for i in range(1000):
+for i in range(horizon):
     actu = (0.0, 0.0, 0.0, 1.0 * (0.5 + math.sin(0.005 * math.pi * i)))
     worm.set_actuation(actu)
     scene.step()

examples/tutorials/batched_IK.py

@@ -1,5 +1,5 @@
+import os
 import numpy as np
-
 import genesis as gs
 
 ########################## init ##########################
@@ -34,10 +34,10 @@
 center = np.tile(np.array([0.4, -0.2, 0.25]), [n_envs, 1])
 angular_speed = np.random.uniform(-10, 10, n_envs)
 r = 0.1
-
 ee_link = robot.get_link("hand")
 
-for i in range(0, 1000):
+horizon = 1000 if "PYTEST_VERSION" not in os.environ else 5
+for i in range(horizon):
     target_pos = np.zeros([n_envs, 3])
     target_pos[:, 0] = center[:, 0] + np.cos(i / 360 * np.pi * angular_speed) * r
     target_pos[:, 1] = center[:, 1] + np.sin(i / 360 * np.pi * angular_speed) * r