robosim.systems

 1"""
 2systems.py - Generic systems in robosim. Should only use entities and components and low level utilities.
 3Import level: 3
 4"""
 5import numpy as np
 6from microecs import World
 7
 8from robosim.camera import Camera
 9from robosim.utils import Point3D
10from robosim.physics import motion_level1, motion_level2, make_grid_3d, check_collision, integrate_velocity_into_pose
11from robosim.components import (HasMotionInput, HasPose, HasCollision, HasVelocity6DoF, HasAcceleration6DoF,
12                                HasPhysicsLevel1, HasPhysicsLevel2, HasModel, HasFPV)
13
14class FPVCameraSystem:
15    """FPV camera system. Basically just sets the position of the camera at each frame"""
16    def __call__(self, world: World):
17        qr = world.query(HasPose, HasFPV)
18        cam: Camera
19        for cam, pose in zip(map(lambda x: x.item(), qr.fpv_camera), qr.pose):
20            cam.set_position(position=pose[0:3, 3], up=pose[0:3, 1], target=(pose[0:3, 3] + pose[0:3, 2]))
21
22class PhysicsSystem:
23    """
24    A orchestration physics system that runs each subtick. It does 4 steps atomically:
25    - candidate motion: move robots based on their physics 'level' (+other forces)
26    - detect collisions and other similar physical events on the candidate new pose
27    - update the candidate position based on the collisions
28    - commit the new physics state
29    """
30
31    def _motion_from_inputs(self, world: World, dt: float, physics_levels: list[str] | None = None):
32        physics_levels = physics_levels or ["level1", "level2"]
33        for physics_level in physics_levels:
34            if physics_level == "level1":
35                qr = world.query(HasPhysicsLevel1, HasMotionInput, HasVelocity6DoF)
36            else: # level 2
37                qr = world.query(HasPhysicsLevel2, HasMotionInput, HasVelocity6DoF, HasAcceleration6DoF)
38
39            if len(qr) == 0:
40                continue
41
42            if physics_level == "level1":
43                vel = motion_level1(motion_input=qr.motion_input, max_velocities=qr.max_velocities)
44                qr.candidate_velocity[:] = vel
45            else: # level 2
46                vel, acc = motion_level2(motion_input=qr.motion_input, dt=dt, current_velocity=qr.velocity,
47                                         max_accelerations=qr.max_accelerations, drag_coefficient=qr.drag_coefficient)
48                qr.candidate_velocity[:] = vel
49                qr.candidate_acceleration[:] = acc
50
51    def _integrate_velocity_into_pose(self, world: World, dt: float):
52        qr = world.query(HasPose, HasVelocity6DoF)
53        qr.candidate_pose[:] = integrate_velocity_into_pose(qr.pose, qr.candidate_velocity, dt)
54
55    def _detect_collisions(self, world: World, cell_size: Point3D):
56        qr = world.query(HasModel, HasCollision, HasPose)
57
58        # Step 1) Update the static collider model using the future pose
59        qr.collider_radii[:] = qr.model_radius * qr.scale
60        qr.collider_bbox[:] = qr.candidate_pose[:, 0:3, 3][:, None] + qr.model_bbox * qr.scale[..., None]
61
62        entities = [world.get_entity(eid) for eid in qr.entity_ids]
63        grid = make_grid_3d(entities, cell_size)
64
65        qr.is_colliding[:] = False
66        for cell_entities in grid.values():
67            for i, e1 in enumerate(cell_entities):
68                for e2 in cell_entities[i+1:]:
69                    if check_collision(e1, e2):
70                        e1.is_colliding = e2.is_colliding = True
71
72    def _update_motion_from_collisions(self, world: World):
73        qr = world.query(HasPose, HasCollision, HasVelocity6DoF)
74        # NOTE: bounce back by negating velocity. We should use some sort of impulse addition + normals
75        # here to figure out the direction of the velocity.
76        qr.candidate_pose[:] = np.where(qr.is_colliding[..., None], qr.pose, qr.candidate_pose)
77        qr.candidate_velocity[:] = np.where(qr.is_colliding, -qr.candidate_velocity, qr.candidate_velocity)
78
79    def _commit(self, world: World):
80        qr = world.query(HasPose, HasVelocity6DoF)
81        qr.pose[:] = qr.candidate_pose
82        qr.velocity[:] = qr.candidate_velocity
83        qr = world.query(HasAcceleration6DoF)
84        qr.acceleration[:] = qr.candidate_acceleration
85
86    def __call__(self, world: World, dt: float, cell_size: Point3D, physics_levels: list[str] | None = None):
87        self._motion_from_inputs(world, dt, physics_levels)
88        self._integrate_velocity_into_pose(world, dt)
89        self._detect_collisions(world, cell_size)
90        self._update_motion_from_collisions(world)
91        self._commit(world)