robolib.systems.physics_system

 1"""physics_system.py - Generic physics in robosim. Should only use entities and components and low level utilities."""
 2import numpy as np
 3from microecs import World
 4
 5from robolib.utils import Point3D, logger
 6from robolib.physics import motion_level1, motion_level2, make_grid_3d, check_collision, integrate_velocity_into_pose
 7from robolib.components import HasMotion, HasPose, HasCollision, HasModel, PhysicsLevel
 8
 9class PhysicsSystem:
10    """
11    A orchestration physics system that runs each subtick. It does 4 steps atomically:
12    - candidate motion: move robots based on their physics 'level' (+other forces)
13    - detect collisions and other similar physical events on the candidate new pose
14    - update the candidate position based on the collisions
15    - commit the new physics state
16    """
17
18    def _motion_from_inputs(self, world: World, dt: float, physics_levels: list[PhysicsLevel] | None = None):
19        qr = world.query(HasMotion)
20        physics_levels = physics_levels or [PhysicsLevel.LEVEL_1, PhysicsLevel.LEVEL_2]
21
22        for level in physics_levels:
23            # TODO: we definitely need a syntax for this pattern (equivalent of: SELECT x where y > 5)
24            mask = (qr.physics_level[:, 0] == level) # (N, ) mask for all the entities (equivalent of where level=level)
25            motion_input = qr.motion_input[mask]
26            if len(motion_input) == 0:
27                continue
28
29            if level == PhysicsLevel.LEVEL_1:
30                vel = motion_level1(motion_input=motion_input, max_velocities=qr.max_velocities[mask])
31                qr.candidate_velocity[mask] = vel
32            elif level == PhysicsLevel.LEVEL_2:
33                vel, acc = motion_level2(motion_input=motion_input, dt=dt, current_velocity=qr.velocity[mask],
34                                         max_accelerations=qr.max_accelerations[mask],
35                                         drag_coefficient=qr.drag_coefficient[mask])
36                qr.candidate_velocity[mask] = vel
37                qr.candidate_acceleration[mask] = acc
38            else:
39                logger.log_every_s(f"Unknown physics level: {level}", "ERROR", freq_s=2)
40
41    def _integrate_velocity_into_pose(self, world: World, dt: float):
42        qr = world.query(HasMotion, HasPose)
43        qr.candidate_pose = integrate_velocity_into_pose(qr.pose, qr.candidate_velocity, dt)
44
45    def _detect_collisions(self, world: World, cell_size: Point3D):
46        qr = world.query(HasModel, HasCollision, HasPose)
47
48        # Step 1) Update the static collider model using the future pose
49        qr.collider_radii = qr.model_radius * qr.scale
50        qr.collider_bbox = qr.candidate_pose[:, 0:3, 3][:, None] + qr.model_bbox * qr.scale[..., None]
51
52        qr.is_colliding = False # note: this broadcasts over all entities in the query result!
53        entities = [world.get_entity(eid) for eid in qr.entity_ids]
54        grid = make_grid_3d(entities, cell_size)
55
56        for cell_entities in grid.values():
57            for i, e1 in enumerate(cell_entities):
58                for e2 in cell_entities[i+1:]:
59                    if check_collision(e1, e2):
60                        e1.is_colliding = e2.is_colliding = True
61
62    def _update_motion_from_collisions(self, world: World):
63        qr = world.query(HasPose, HasCollision, HasMotion)
64        # NOTE: bounce back by negating velocity. We should use some sort of impulse addition + normals
65        # here to figure out the direction of the velocity.
66        qr.candidate_pose = np.where(qr.is_colliding[..., None], qr.pose, qr.candidate_pose)
67        qr.candidate_velocity = np.where(qr.is_colliding, -qr.candidate_velocity, qr.candidate_velocity)
68
69    def _commit(self, world: World):
70        qr = world.query(HasPose, HasMotion)
71        qr.pose = qr.candidate_pose
72        qr.velocity = qr.candidate_velocity
73        qr.acceleration = qr.candidate_acceleration
74
75    def __call__(self, world: World, dt: float, cell_size: Point3D, physics_levels: list[PhysicsLevel] | None = None):
76        self._motion_from_inputs(world, dt, physics_levels)
77        self._integrate_velocity_into_pose(world, dt)
78        self._detect_collisions(world, cell_size)
79        self._update_motion_from_collisions(world)
80        self._commit(world)