urdf loaded obstacles pushed

Author: saiccoumar

simulator_testing/collision_checking/mujoco_collision.py

@@ -1,50 +1,88 @@
 #!/usr/bin/env python3
 """
 This script benchmarks collision checking for two MuJoCo models by randomizing joint positions within the URDF-defined joint limits.
-It also adds static obstacles and checks for collisions both within the robot (self-collisions) and between the robot and obstacles.
+It also adds static obstacles from a URDF file and checks for collisions both within the robot (self-collisions) and between the robot and obstacles.
 Additionally, an interactive visualization mode is provided that uses MuJoCo's viewer.
 """
-import csv
+
 import mujoco
 import mujoco.viewer
 import numpy as np
 import time
 import sys
 from scipy.spatial.transform import Rotation
 
-def create_static_obstacles(spec, num_obstacles=100):
-    obstacles = []
-    obstacle_geom_ids = []
-    
-    for i in range(num_obstacles):
-        size = np.random.uniform(0.05, 0.2)
-        pos = np.random.uniform(-2, 2, size=3)
-        
-        euler_angles = np.random.uniform(-np.pi, np.pi, size=3)
-        rotation = Rotation.from_euler('xyz', euler_angles)
-        quat = rotation.as_quat()
-        
-        body = spec.worldbody.add_body(
-            name=f"box{i}",
-            pos=pos,
-            quat=quat
-        )
-        geom = body.add_geom(
-            type=mujoco.mjtGeom.mjGEOM_BOX,
-            size=[size, size, size],
-            rgba=[1, 0, 0, 1],
-            condim=3,
-            friction=[1, 0.1, 0.1]
-        )
-        
-        obstacles.append(body)
-        obstacle_geom_ids.append(geom)
-    
-    return obstacles, obstacle_geom_ids
 
+spec_a = mujoco.MjSpec.from_file("../../assets/panda/panda.urdf")
+spec_b = mujoco.MjSpec.from_file("../../assets/panda/spherized_panda.urdf")
+
+# Add a ground plane to each model.
+spec_a.worldbody.add_geom(
+    type=mujoco.mjtGeom.mjGEOM_PLANE,
+    size=[2, 2, 0.1],
+    rgba=[0.2, 0.2, 0.2, 1]
+)
 
 
+spec_b.worldbody.add_geom(
+    type=mujoco.mjtGeom.mjGEOM_PLANE,
+    size=[2, 2, 0.1],
+    rgba=[0.2, 0.2, 0.2, 1]
+)
 
+# Load obstacles from obstacles.urdf
+def load_obstacles_from_urdf(spec, urdf_path):
+    """
+    Loads obstacles from a given URDF file and adds them to the simulation.
+    Returns a list of added geometry objects.
+    """
+    obstacle_geoms = []
+    
+    # Load the URDF file
+    obstacle_spec = mujoco.MjSpec.from_file(urdf_path)
+
+    for geom in obstacle_spec.worldbody.geoms:  # Changed from geom to geoms
+        obstacle_geom = spec.worldbody.add_geom(
+            type=geom.type,
+            size=geom.size,
+            pos=geom.pos if hasattr(geom, 'pos') else None,
+            quat=geom.quat if hasattr(geom, 'quat') else None,
+            rgba=geom.rgba if hasattr(geom, 'rgba') else [1, 1, 1, 1],
+            condim=geom.condim if hasattr(geom, 'condim') else 3,
+            friction=geom.friction if hasattr(geom, 'friction') else [1, 0.1, 0.1]
+        )
+        obstacle_geoms.append(obstacle_geom)
+    
+    # Handle any child bodies recursively
+    if hasattr(obstacle_spec.worldbody, 'body'):
+        for body in obstacle_spec.worldbody.body:
+            # Add a new body to the spec
+            new_body = spec.worldbody.add_body()
+            
+            # Copy over the body properties
+            if hasattr(body, 'pos'):
+                new_body.pos = body.pos
+            if hasattr(body, 'quat'):
+                new_body.quat = body.quat
+                
+            # Add all geoms from this body
+            for geom in body.geoms:  # Changed from geom to geoms
+                obstacle_geom = new_body.add_geom(
+                    type=geom.type,
+                    size=geom.size,
+                    pos=geom.pos if hasattr(geom, 'pos') else None,
+                    quat=geom.quat if hasattr(geom, 'quat') else None,
+                    rgba=geom.rgba if hasattr(geom, 'rgba') else [1, 1, 1, 1],
+                    condim=geom.condim if hasattr(geom, 'condim') else 3,
+                    friction=geom.friction if hasattr(geom, 'friction') else [1, 0.1, 0.1]
+                )
+                obstacle_geoms.append(obstacle_geom)
+    return obstacle_geoms
+
+obstacles_urdf_path = "../obstacles.urdf"
+obstacles_geoms = load_obstacles_from_urdf(spec_a, obstacles_urdf_path)
+
+# Function to randomize joint positions within the limits of the URDF
 def random_joint_positions_from_limits(model):
     """
     Generates random joint positions that respect the joint limits provided
@@ -54,7 +92,6 @@ def random_joint_positions_from_limits(model):
     qpos = np.empty(model.nq)
     for i in range(model.njnt):
         lower, upper = model.jnt_range[i]
-        # If the joint is fixed (limits are equal), use that value.
         if lower == upper:
             qpos[i] = lower
         else:
@@ -76,10 +113,6 @@ def debug_simulation(model, data, step, threshold=1e3):
         print("---------------------------")
 
 def benchmark_collision_checking(spec, obstacle_geom_ids, num_samples=100_000):
-    """
-    Evaluate average self-collision counts, obstacle collision counts,
-    elapsed time, and collision check total time over multiple samples.
-    """
     model = spec.compile()
     data = mujoco.MjData(model)
 
@@ -90,9 +123,6 @@ def benchmark_collision_checking(spec, obstacle_geom_ids, num_samples=100_000):
     for geom_obj in obstacle_geom_ids:
         obstacle_geom_indices.add(geom_obj.id)
 
-    print("Obstacle geom indices:", obstacle_geom_indices)
-    # print("Robot geom indices:", robot_geom_indices)
-
     self_collision_count = 0
     obstacle_collision_count = 0
     collision_check_total_time = 0.0
@@ -130,27 +160,19 @@ def benchmark_collision_checking(spec, obstacle_geom_ids, num_samples=100_000):
 
 
 def benchmark_collision_checking_visual(spec, obstacle_geom_ids, sleep_time=0.0001, num_samples=100_000):
-    """
-    Evaluate average self-collision counts, obstacle collision counts,
-    elapsed time, and collision check total time over multiple samples.
-    """
     model = spec.compile()
     data = mujoco.MjData(model)
-    print("Model nq =", model.nq, "model njnt =", model.njnt)
     obstacle_geom_indices = set()
     for geom_obj in obstacle_geom_ids:
         obstacle_geom_indices.add(geom_obj.id)
 
-    print("Obstacle geom indices:", obstacle_geom_indices)
     self_collision_count = 0
     obstacle_collision_count = 0
     collision_check_total_time = 0.0
 
     start_time = time.perf_counter()
 
     with mujoco.viewer.launch_passive(model, data) as viewer:
-        print("Starting static visualization. Close the window to exit.")
-        
         for _ in range(num_samples):
             data.qpos[:] = random_joint_positions_from_limits(model)
             data.qvel[:] = 0
@@ -168,8 +190,6 @@ def benchmark_collision_checking_visual(spec, obstacle_geom_ids, sleep_time=0.00
                         obstacle_collision_count += 1
                     else:
                         self_collision_count += 1
-                # break
-            
             viewer.sync()
             collision_check_total_time += time.perf_counter() - start_check
             time.sleep(sleep_time)
@@ -183,132 +203,26 @@ def benchmark_collision_checking_visual(spec, obstacle_geom_ids, sleep_time=0.00
 
     return avg_self_collisions, avg_obstacle_collisions, avg_elapsed_time, avg_collision_check_time
 
-# if __name__ == "__main__":
-#     spec_a = mujoco.MjSpec.from_file("../../assets/panda/panda.urdf")
-#     spec_b = mujoco.MjSpec.from_file("../../assets/panda/spherized_panda.urdf")
-
-#     # Add a ground plane to each model.
-#     spec_a.worldbody.add_geom(
-#         type=mujoco.mjtGeom.mjGEOM_PLANE,
-#         size=[2, 2, 0.1],
-#         rgba=[0.2, 0.2, 0.2, 1]
-#     )
-
-
-#     spec_b.worldbody.add_geom(
-#         type=mujoco.mjtGeom.mjGEOM_PLANE,
-#         size=[2, 2, 0.1],
-#         rgba=[0.2, 0.2, 0.2, 1]
-#     )
-
-
-#     # Create obstacles for both models.
-#     obstacles_a, obstacles_a_geoms = create_static_obstacles(spec_a, 100)
-#     obstacles_b, obstacles_b_geoms = create_static_obstacles(spec_b, 100)
-#     # If '--visualize' is passed as a command-line argument, run the visualization mode.
-#     if '--visualize' in sys.argv:
-#         # Run the benchmark (this is your original benchmarking code)
-#         print("Benchmarking Model A (Complex)")
-#         a_self_count, a_obst_count, a_total_time, a_check_time = benchmark_collision_checking_visual(spec_a, obstacles_a_geoms, sleep_time=0)
-#         print("Benchmarking Model B (Simple)")
-#         b_self_count, b_obst_count, b_total_time, b_check_time = benchmark_collision_checking_visual(spec_b, obstacles_b_geoms, sleep_time=0)
-#     else:
-#         # Run the benchmark (this is your original benchmarking code)
-#         print("Benchmarking Model A (Complex)")
-#         a_self_count, a_obst_count, a_total_time, a_check_time = benchmark_collision_checking(spec_a, obstacles_a_geoms)
-#         print("Benchmarking Model B (Simple)")
-#         b_self_count, b_obst_count, b_total_time, b_check_time = benchmark_collision_checking(spec_b, obstacles_b_geoms)
-
-#     print(f"\nModel A (Complex):")
-#     print(f"  Average Simulation Time: {a_total_time:.3f} sec")
-#     print(f"  Average Self-Collisions: {a_self_count}")
-#     print(f"  Average Obstacle Collisions: {a_obst_count}")
-#     print(f"  Average Collision Check Time: {a_check_time:.6f} sec")
-    
-#     print(f"\nModel B (Simple):")
-#     print(f"  Average Simulation Time: {b_total_time:.3f} sec")
-#     print(f"  Average Self-Collisions: {b_self_count}")
-#     print(f"  Average Obstacle Collisions: {b_obst_count}")
-#     print(f"  Average Collision Check Time: {b_check_time:.6f} sec")
-
 if __name__ == "__main__":
-    # Loop over the specified indices.
-    results_summary = []
-
-    for i in range(1):
-        urdf_file = f"../../assets/panda/panda.urdf"
-        print(f"\nBenchmarking {urdf_file}")
-
-        # Load the model from the current URDF file.
-        spec = mujoco.MjSpec.from_file(urdf_file)
-
-        # Add a ground plane to the model.
-        spec.worldbody.add_geom(
-            type=mujoco.mjtGeom.mjGEOM_PLANE,
-            size=[2, 2, 0.1],
-            rgba=[0.2, 0.2, 0.2, 1]
-        )
-
-        # Create static obstacles for the current model.
-        obstacles, obstacles_geoms = create_static_obstacles(spec, 100)
-
-        # Run benchmark in either visualization mode or headless.
-        if '--visualize' in sys.argv:
-            self_count, obst_count, total_time, check_time = benchmark_collision_checking_visual(
-                spec, obstacles_geoms, sleep_time=0
-            )
-        else:
-            self_count, obst_count, total_time, check_time = benchmark_collision_checking(
-                spec, obstacles_geoms
-            )
-
-        # Save the results in the summary list.
-        # Save the results in the summary list, including the sphere count.
-        results_summary.append({
-            "sphere_count": i,
-            "urdf_file": urdf_file,
-            "self_count": self_count,
-            "obst_count": obst_count,
-            "total_time": total_time,
-            "check_time": check_time
-        })
-
-        # Print the benchmark results for this file.
-        print(f"Results for {urdf_file}:")
-        print(f"  Average Simulation Time: {total_time:.3f} sec")
-        print(f"  Average Self-Collisions: {self_count}")
-        print(f"  Average Obstacle Collisions: {obst_count}")
-        print(f"  Average Collision Check Time: {check_time:.6f} sec")
-
-    # Print a summary of all the results after processing all files.
-    print("\nSummary of all results:")
-    for result in results_summary:
-        print(f"File: {result['urdf_file']}")
-        print(f"  Sphere Count: {result['sphere_count']}")
-        print(f"  Average Simulation Time: {result['total_time']:.3f} sec")
-        print(f"  Average Self-Collisions: {result['self_count']}")
-        print(f"  Average Obstacle Collisions: {result['obst_count']}")
-        print(f"  Average Collision Check Time: {result['check_time']:.6f} sec")
-        print("-" * 40)
-
-    csv_filename = "tune_sphere_count_panda.csv"
-    with open(csv_filename, 'a', newline='') as csvfile:
-        fieldnames = [
-            'Sphere count',
-            'Average Simulation Time',
-            'Average Self-Collisions',
-            'Average Obstacle Collisions',
-            'Average Collision Check Time'
-        ]
-        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
-        writer.writeheader()
-        for result in results_summary:
-            writer.writerow({
-                'Sphere count': result['sphere_count'],
-                'Average Simulation Time': result['total_time'],
-                'Average Self-Collisions': result['self_count'],
-                'Average Obstacle Collisions': result['obst_count'],
-                'Average Collision Check Time': result['check_time']
-            })
-
-    print(f"\nBenchmark results saved to {csv_filename}")
+    # Load obstacles from the URDF file and add them to both models
+    obstacles_a_geoms = load_obstacles_from_urdf(spec_a, obstacles_urdf_path)
+    obstacles_b_geoms = load_obstacles_from_urdf(spec_b, obstacles_urdf_path)
+
+    # Run the benchmark for both models
+    print("Benchmarking Model A (Complex)")
+    a_self_count, a_obst_count, a_total_time, a_check_time = benchmark_collision_checking(spec_a, obstacles_a_geoms)
+    print("Benchmarking Model B (Simple)")
+    b_self_count, b_obst_count, b_total_time, b_check_time = benchmark_collision_checking(spec_b, obstacles_b_geoms)
+
+    # Output the results
+    print(f"\nModel A (Complex):")
+    print(f"  Average Simulation Time: {a_total_time:.3f} sec")
+    print(f"  Average Self-Collisions: {a_self_count}")
+    print(f"  Average Obstacle Collisions: {a_obst_count}")
+    print(f"  Average Collision Check Time: {a_check_time:.6f} sec")
+    
+    print(f"\nModel B (Simple):")
+    print(f"  Average Simulation Time: {b_total_time:.3f} sec")
+    print(f"  Average Self-Collisions: {b_self_count}")
+    print(f"  Average Obstacle Collisions: {b_obst_count}")
+    print(f"  Average Collision Check Time: {b_check_time:.6f} sec")