robolib.systems

13class RenderSystem:
14    """Only the main 3D render camera(s). For FPV rendering on a texture, use FPVRenderSystem."""
15    def __call__(self, world: World, camera: Camera, extra_drawables: list[Drawable] | None = None,
16                 camera_entity_id: EntityId | None = None, wireframe_mode: bool = False,
17                 collision_mode: bool = False, display_uav_trace: bool = False, uav_traces: dict | None = None,
18                 display_hud: bool = False, extra_callables: list[Callable] | None = None):
19        rl.ClearBackground(DEFAULT_BACKGROUND)
20
21        rl.BeginMode3D(camera.camera[0])
22
23        if wireframe_mode:
24            rl.rlEnableWireMode()
25
26        if collision_mode:
27            qr = world.query(HasModel, HasPose, exclude=[HasCollision])
28            for model, pose, scale in zip(qr.model, qr.pose, qr.scale):
29                rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
30
31            qr = world.query(HasModel, HasPose, GroundFloor)
32            for model, pose, scale in zip(qr.model, qr.pose, qr.scale):
33                rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
34
35            qr = world.query(HasModel, HasPose, HasCollision, exclude=[GroundFloor])
36            collider_size = qr.collider_bbox[:, 1] - qr.collider_bbox[:, 0]
37            for pose, kind, radius, size, colliding in zip(qr.pose, qr.collider_kind, qr.collider_radii,
38                                                           collider_size, qr.is_colliding):
39                color = rl.RED if colliding else rl.GREEN
40                if kind == ColliderKinds.SPHERE:
41                    rl.DrawSphere(pose[0:3, 3].tolist(), radius.item(), color)
42                else: # AABB
43                    rl.DrawCube(pose[0:3, 3].tolist(), *size.tolist(), color)
44
45        else:
46            qr = world.query(HasModel, HasPose)
47            for i, (model, pose, scale) in enumerate(zip(qr.model, qr.pose, qr.scale)):
48                if camera_entity_id is not None and qr.entity_ids[i] == camera_entity_id: # don't draw yourself in FPV
49                    continue
50                rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
51
52        for drawable in (extra_drawables or []): # e.g. plugins
53            drawable.draw(world=world, entity_id=camera_entity_id)
54        for callable_fn in (extra_callables or []):
55            callable_fn(world=world)
56
57        if display_uav_trace:
58            assert uav_traces is not None, uav_traces
59            for traces in uav_traces.values():
60                for trace in traces.keys():
61                    rl.DrawSphere(trace, 0.05, rl.GREEN)
62
63        rl.EndMode3D()
64
65        if wireframe_mode:
66            rl.rlDisableWireMode()
67
68        if display_hud:
69            n_robots = len(world.query(HasFPV))
70            msg = f"{camera.label}\nObjects: {len(world)} (robots: {n_robots})"
71            rl.DrawText(msg.encode(), 10, 10, 20, rl.DARKGRAY)
72            rl.DrawFPS(rl.GetScreenWidth() - 80, 10)

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

12class FPVRenderSystem:
13    """FPVRenderSystem - renders all the connected eids to custom textures that are streamed to the network clients"""
14    def __call__(self, world: World, connected_eids: list[EntityId], extra_drawables: list[Drawable] | None = None):
15        for eid in connected_eids:
16            entity = world.get_entity(eid)
17            fpv_texture: "rl.RenderTexture" = entity.fpv_texture.item()
18            fpv_data: FPVData = entity.fpv_data.item()
19            camera: Camera = entity.fpv_camera.item()
20
21            rl.BeginTextureMode(fpv_texture)
22            rl.ClearBackground(DEFAULT_BACKGROUND)
23            rl.BeginMode3D(camera.camera[0])
24
25            qr = world.query(HasModel, HasPose)
26            for i, (model, pose, scale) in enumerate(zip(qr.model, qr.pose, qr.scale)):
27                if qr.entity_ids[i] == eid: # don't draw yourself in FPV
28                    continue
29                rl.DrawModel(model.item().model, pose[0:3, 3].tolist(), scale.item(), rl.WHITE)
30
31            for drawable in (extra_drawables or []): # e.g. plugins
32                drawable.draw(world, entity_id=eid)
33
34            rl.EndMode3D()
35            rl.EndTextureMode()
36
37            fpv_img = rl.LoadImageFromTexture(fpv_texture.texture)
38            fpv_img_ptr = rl.ffi.new("Image *", fpv_img)
39            rl.ImageFlipVertical(fpv_img_ptr)
40            with fpv_data.lock:
41                fpv_data.frame = bytes(rl.ffi.buffer(fpv_img_ptr.data, len(fpv_data.frame)))
42                fpv_data.frame_id += 1
43                fpv_data._frame_compressed = None # pylint: disable=protected-access
44            rl.UnloadImage(fpv_img_ptr[0])