Agent2Robot / src /simulation /physics_simulator.py
sam133
Fix physics simulator implementation and imports
7c1d4f6
Raw
History Blame Contribute Delete
3 kB
"""
Physics simulation module for robot designs.
"""
import pybullet as p
import numpy as np
from typing import Dict, Any, List, Tuple
from src.core.robot_design import RobotDesign
class PhysicsSimulator:
"""Base class for physics simulation."""
def __init__(self):
"""Initialize the physics simulator."""
self.client = p.connect(p.DIRECT)
p.setGravity(0, 0, -9.81)
def simulate(self, design: RobotDesign) -> Dict[str, Any]:
"""Simulate the robot design and return performance metrics."""
raise NotImplementedError("Subclasses must implement simulate()")
def cleanup(self):
"""Clean up simulation resources."""
p.disconnect(self.client)
class PyBulletSimulator(PhysicsSimulator):
"""PyBullet-based physics simulator implementation."""
def simulate(self, design: RobotDesign) -> Dict[str, Any]:
"""Simulate the robot design using PyBullet."""
try:
# Create ground plane
p.createMultiBody(
baseMass=0,
baseCollisionShapeIndex=p.createCollisionShape(p.GEOM_PLANE),
basePosition=[0, 0, 0]
)
# Create robot body
robot_id = p.createMultiBody(
baseMass=design.mass,
baseCollisionShapeIndex=p.createCollisionShape(
p.GEOM_BOX,
halfExtents=[design.dimensions[0]/2, design.dimensions[1]/2, design.dimensions[2]/2]
),
basePosition=[0, 0, design.dimensions[2]/2]
)
# Run simulation
results = {
"stability": self._check_stability(robot_id),
"performance": self._measure_performance(robot_id, design)
}
return results
except Exception as e:
return {"error": str(e)}
finally:
self.cleanup()
def _check_stability(self, robot_id: int) -> float:
"""Check robot stability during simulation."""
# Simple stability check based on final position
final_pos, _ = p.getBasePositionAndOrientation(robot_id)
return 1.0 if abs(final_pos[2]) < 0.1 else 0.0
def _measure_performance(self, robot_id: int, design: RobotDesign) -> Dict[str, float]:
"""Measure robot performance metrics."""
return {
"speed": design.max_speed,
"efficiency": design.efficiency,
"stability": self._check_stability(robot_id)
}
class SimulationFactory:
"""Factory for creating physics simulators."""
@staticmethod
def create_simulator(simulator_type: str = "pybullet") -> PhysicsSimulator:
"""Create a physics simulator instance."""
if simulator_type == "pybullet":
return PyBulletSimulator()
else:
raise ValueError(f"Unsupported simulator type: {simulator_type}")