update

Author: Kashu7100

examples/rigid/grasp_bottle.py

@@ -8,6 +8,7 @@
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("-v", "--vis", action="store_true", default=False)
+    parser.add_argument("-n", "--n_envs", type=int, default=49)
     args = parser.parse_args()
 
     ########################## init ##########################
@@ -48,12 +49,16 @@ def main():
     )
 
     ########################## build ##########################
-    scene.build()
+    scene.build(n_envs=args.n_envs, env_spacing=(1,1))
 
     motors_dof = np.arange(7)
     fingers_dof = np.arange(7, 9)
 
     # Optional: set control gains
+    if args.n_envs == 0:
+        franka.set_qpos(np.array([1.56, -0.72, -0.02, -2.09, 0.04, 1.33, 2.4, 0.01, 0.01]))
+    else:
+        franka.set_qpos(np.array([[1.56, -0.72, -0.02, -2.09, 0.04, 1.33, 2.4, 0.01, 0.01]] * args.n_envs))
     franka.set_dofs_kp(
         np.array([4500, 4500, 3500, 3500, 2000, 2000, 2000, 100, 100]),
     )
@@ -70,10 +75,11 @@ def main():
     # move to pre-grasp pose
     qpos = franka.inverse_kinematics(
         link=end_effector,
-        pos=np.array([0.65, 0.0, 0.25]),
-        quat=np.array([0, 1, 0, 0]),
+        pos =np.array([0.65, 0.0, 0.25]) if args.n_envs == 0 else np.array([[0.65, 0.0, 0.25]] * args.n_envs),
+        quat=np.array([0, 1, 0, 0]) if args.n_envs == 0 else np.array([[0, 1, 0, 0]] * args.n_envs),
     )
-    qpos[-2:] = 0.04
+    qpos[...,-2:] = 0.04
+
     path = franka.plan_path(qpos)
     for waypoint in path:
         franka.control_dofs_position(waypoint)
@@ -84,30 +90,36 @@ def main():
     # reach
     qpos = franka.inverse_kinematics(
         link=end_effector,
-        pos=np.array([0.65, 0.0, 0.142]),
-        quat=np.array([0, 1, 0, 0]),
+        pos= np.array([0.65, 0.0, 0.142]) if args.n_envs == 0 else np.array([[0.65, 0.0, 0.142]] * args.n_envs),
+        quat=np.array([0, 1, 0, 0]) if args.n_envs == 0 else np.array([[0, 1, 0, 0]] * args.n_envs),
     )
-    franka.control_dofs_position(qpos[:-2], motors_dof)
+    franka.control_dofs_position(qpos[...,:-2], motors_dof)
     for i in range(100):
         scene.step()
 
     # grasp
-    franka.control_dofs_position(qpos[:-2], motors_dof)
-    franka.control_dofs_position(np.array([0, 0]), fingers_dof)  # you can use position control
+    franka.control_dofs_position(qpos[...,:-2], motors_dof)
+    franka.control_dofs_position(
+        np.array([0, 0]) if args.n_envs == 0 else np.array([[0, 0]] * args.n_envs), 
+        fingers_dof
+    )  # you can use position control
     for i in range(100):
         scene.step()
 
     # lift
     qpos = franka.inverse_kinematics(
         link=end_effector,
-        pos=np.array([0.65, 0.0, 0.3]),
-        quat=np.array([0, 1, 0, 0]),
+        pos=np.array([0.65, 0.0, 0.3]) if args.n_envs == 0 else np.array([[0.65, 0.0, 0.3]] * args.n_envs),
+        quat=np.array([0, 1, 0, 0]) if args.n_envs == 0 else np.array([[0, 1, 0, 0]] * args.n_envs),
     )
-    franka.control_dofs_position(qpos[:-2], motors_dof)
-    franka.control_dofs_force(np.array([-20, -20]), fingers_dof)  # can also use force control
+    franka.control_dofs_position(qpos[...,:-2], motors_dof)
+    franka.control_dofs_force(
+        np.array([-20, -20]) if args.n_envs == 0 else np.array([[-20, -20]] * args.n_envs),
+        fingers_dof
+    )  # can also use force control
     for i in range(1000):
         scene.step()
 
 
 if __name__ == "__main__":
-    main()
+    main()