updated mujoco cc with urdf blocks

Author: NiharKod

simulator_testing/collision_checking/mujoco_collision_new.py

@@ -0,0 +1,228 @@
+#!/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 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 mujoco
+import mujoco.viewer
+import numpy as np
+import time
+import sys
+from scipy.spatial.transform import Rotation
+
+
+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
+    in the URDF for a fixed-base robot. Assumes that model.nq (number of generalized
+    coordinates) equals model.njnt (number of joints with limits).
+    """
+    qpos = np.empty(model.nq)
+    for i in range(model.njnt):
+        lower, upper = model.jnt_range[i]
+        if lower == upper:
+            qpos[i] = lower
+        else:
+            qpos[i] = np.random.uniform(lower, upper)
+    return qpos
+
+def debug_simulation(model, data, step, threshold=1e3):
+    unstable = False
+    for i, acc in enumerate(data.qacc):
+        if np.isnan(acc) or np.isinf(acc) or np.abs(acc) > threshold:
+            print(f"Warning: Unstable acceleration at DOF {i} on step {step}: {acc}")
+            unstable = True
+    if unstable:
+        print("---- Debug Information ----")
+        print(f"Time: {data.time:.4f}")
+        print("qpos:", data.qpos)
+        print("qvel:", data.qvel)
+        print("qacc:", data.qacc)
+        print("---------------------------")
+
+def benchmark_collision_checking(spec, obstacle_geom_ids, num_samples=100_000):
+    model = spec.compile()
+    data = mujoco.MjData(model)
+    
+    print(model.ngeom)
+    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)
+
+    self_collision_count = 0
+    obstacle_collision_count = 0
+    collision_check_total_time = 0.0
+
+    start_time = time.perf_counter()
+
+    for _ in range(num_samples):
+        data.qpos[:] = random_joint_positions_from_limits(model)
+        data.qvel[:] = 0
+        data.ctrl[:] = 0 
+
+        start_check = time.perf_counter()
+        mujoco.mj_forward(model, data)
+        
+        for i in range(data.ncon):
+            contact = data.contact[i]
+            if contact.dist < 0:  # Collision detected
+                geom1, geom2 = contact.geom1, contact.geom2
+
+                if geom1 in obstacle_geom_indices or geom2 in obstacle_geom_indices:
+                    obstacle_collision_count += 1
+                else:
+                    self_collision_count += 1
+
+        collision_check_total_time += time.perf_counter() - start_check
+
+    elapsed_time = time.perf_counter() - start_time
+    
+    avg_self_collisions = self_collision_count / num_samples
+    avg_obstacle_collisions = obstacle_collision_count / num_samples
+    avg_collision_check_time = collision_check_total_time / num_samples
+    avg_elapsed_time = elapsed_time / num_samples
+    
+    return avg_self_collisions, avg_obstacle_collisions, avg_elapsed_time, avg_collision_check_time
+
+
+def benchmark_collision_checking_visual(spec, obstacle_geom_ids, sleep_time=0.0001, num_samples=100_000):
+    model = spec.compile()
+    data = mujoco.MjData(model)
+    obstacle_geom_indices = set()
+    for geom_obj in obstacle_geom_ids:
+        obstacle_geom_indices.add(geom_obj.id)
+
+    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:
+        for _ in range(num_samples):
+            data.qpos[:] = random_joint_positions_from_limits(model)
+            data.qvel[:] = 0
+            data.ctrl[:] = 0 
+
+            start_check = time.perf_counter()
+            mujoco.mj_forward(model, data)
+            
+            for i in range(data.ncon):
+                contact = data.contact[i]
+                if contact.dist < 0:  # Collision detected
+                    geom1, geom2 = contact.geom1, contact.geom2
+
+                    if geom1 in obstacle_geom_indices or geom2 in obstacle_geom_indices:
+                        obstacle_collision_count += 1
+                    else:
+                        self_collision_count += 1
+            viewer.sync()
+            collision_check_total_time += time.perf_counter() - start_check
+            time.sleep(sleep_time)
+
+    elapsed_time = time.perf_counter() - start_time
+    
+    avg_self_collisions = self_collision_count / num_samples
+    avg_obstacle_collisions = obstacle_collision_count / num_samples
+    avg_collision_check_time = collision_check_total_time / num_samples
+    avg_elapsed_time = elapsed_time / num_samples
+
+    return avg_self_collisions, avg_obstacle_collisions, avg_elapsed_time, avg_collision_check_time
+
+if __name__ == "__main__":
+    # 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")