HyperDrone

%pip install -q "hyperdrone[examples]" foundation-policy==1.0.1

Show imports and helpers

Hide imports and helpers

from itertools import cycle
from contextlib import closing
from dataclasses import replace
from pathlib import Path

import imageio.v2 as imageio
import matplotlib
import numpy as np
from foundation_policy import Raptor
from IPython.display import Video, display
from scipy.spatial.transform import Rotation
from PIL import Image, ImageDraw, ImageFont, ImageOps

import hyperlux
from hyperlux.camera import CameraSpec, RGB, Depth, Segmentation, Normals, Flow, FlowStatus
from hyperdrone import Config, HyperDrone
from hyperdrone.examples.data import rigged_model_path
from hyperdrone.examples.shots import HYPERDRONE_SHOTS as shots

hyperlux.set_verbosity("warning")


FONT = ImageFont.truetype(str(Path(matplotlib.get_data_path()) / "fonts/ttf/DejaVuSans.ttf"), 16)


def rgb(image):
    return image if image.dtype == np.uint8 else (np.clip(image, 0, 1) * 255).astype(np.uint8)


def writer(path, fps):
    return imageio.get_writer(str(path), fps=fps, codec="libx264", quality=8, macro_block_size=1)


def panel(images, titles):
    images = [rgb(image) for image in images]
    height, width = images[0].shape[:2]
    canvas = Image.new("RGB", (width * len(images), height + 32), "#111827")
    draw = ImageDraw.Draw(canvas)
    for i, (image, title) in enumerate(zip(images, titles)):
        canvas.paste(Image.fromarray(image), (i * width, 32))
        draw.text((i * width + 12, 6), title, font=FONT, fill="#e5e7eb")
    return np.asarray(canvas)


def sensor_panel(depth, labels, normals, flow, status):
    depth_image = matplotlib.colormaps["viridis"](np.clip(depth / 10, 0, 1))[..., :3]
    depth_image[labels == 0xFFFFFFFF] = 0
    label_image = matplotlib.colormaps["tab20"](labels % 20)[..., :3]
    label_image[labels == 0xFFFFFFFF] = 0
    normal_image = np.clip(normals * 0.5 + 0.5, 0, 1)
    normal_image[np.all(normals == 0, axis=-1)] = 0
    flow_image = matplotlib.colormaps["magma"](np.clip(np.linalg.norm(flow, axis=-1) / 50, 0, 1))[..., :3]
    flow_image[status != hyperlux.FlowStatus.VALID] = 0
    images = [np.asarray(Image.fromarray(rgb(image)).resize((WIDTH // 2, HEIGHT // 2), Image.Resampling.NEAREST))
              for image in (depth_image, label_image, normal_image, flow_image)]
    return panel(images, ["FPV depth | 0-10 m", "FPV semantic IDs",
                          "FPV world normals", "FPV flow | 0-50 px"])


def join(paths, output, fps):
    with imageio.get_reader(str(paths[0])) as first:
        size = first.get_meta_data()["size"]
    with writer(output, fps) as out:
        for path in paths:
            with imageio.get_reader(str(path)) as source:
                for frame in source:
                    out.append_data(np.asarray(ImageOps.pad(Image.fromarray(frame), size, color="#111827")))


def show_video(path):
    display(Video(str(path), embed=True, width=768, html_attributes='controls loop style="max-width: 100%; height: auto;"'))


def camera(position, pitch, yaw, fov):
    orientation = Rotation.from_euler("ZY", [yaw, pitch]).as_quat()[[3, 0, 1, 2]].tolist()
    return {"fov": fov, "mount": {"pose": {"position": position, "orientation": orientation}}}


def block(state, layout, name):
    offset, size = layout.blocks[name]
    return state[:, offset:offset + size]

def pilots(env):
    policies = [Raptor() for _ in range(env.total_instances * env.num_agents)]
    for policy in policies:
        policy.reset()
    return policies, np.zeros((env.total_instances, env.action_dim), dtype=np.float32)

def act(env, policies, previous, target=0, layout=None):
    layout = env.observation_layout_privileged if layout is None else layout
    state = env.observe_privileged().reshape(env.total_instances * env.num_agents, -1)
    prefix = "agent0/" if env.num_agents > 1 else ""
    observation = np.concatenate([
        block(state, layout, prefix + "position") - target,
        *[block(state, layout, prefix + name) for name in
          ("orientation_rotation_matrix", "linear_velocity", "angular_velocity")],
        previous.reshape(len(policies), -1),
    ], axis=1)
    return np.concatenate([policy.evaluate_step(row[None]) for policy, row in zip(policies, observation)]).astype(np.float32).reshape(env.total_instances, env.action_dim)

Setup

dataset = hyperlux.datasets.procthor.Train200()
scenes = [dataset[index].reference for index in (1, 0)]
WIDTH, HEIGHT = 640, 360
STEPS, STRIDE, FPS = 500, 4, 25
output = Path("hyperdrone_videos")
output.mkdir(exist_ok=True)
drone_asset = rigged_model_path("x500")
config = Config(instances_per_world=1, width=WIDTH, height=HEIGHT, preset="x500_fpv", motion_blur_samples=8)
initial = {"mdp": {"init": {
    "guidance": 0.0, "max_position": 0.25, "max_angle": 0.1,
    "max_linear_velocity": 0.0, "max_angular_velocity": 0.0,
    "min_rpm": 0.0, "max_rpm": 0.0,
}}}

fpv = camera([0.1, 0, 0.19], 0.3, 0, 100)
chase = camera([-0.8, -0.8, 0.6], 0.45, np.pi / 4, 65)

Fly the preset

sensor_config = replace(config, dataset="procthor", cameras=(CameraSpec("front", (RGB("float32"), Depth(), Segmentation(labels="semantic"), Normals(), Flow(), FlowStatus())), CameraSpec("chase", (RGB("float32"),))))
with closing(HyperDrone(scenes, config=sensor_config, drone_asset=drone_asset,
                              parameters={"dynamics": initial, "cameras": {"front": fpv, "chase": chase}})) as env:
    print("Rendering backend:", env.render_backend)
    env.reset(spawn_points=shots["flight"])
    env.observe()
    semantic_classes = env.segmentation_classes
    semantic_ids = env.outputs[Segmentation(labels="semantic")].read()[0, 0]
    print("Visible semantic IDs:", {int(id): semantic_classes[int(id)]
                                    for id in np.unique(semantic_ids) if id != 0xFFFFFFFF})
    base_layout = env.observation_layout_privileged
    image_layout = env.observation_layout
    policies, previous = pilots(env)
    assert np.isclose(STRIDE * env.dt, 1 / FPS)
    with writer(output / "01-preset.mp4", FPS) as video:
        for step in range(STEPS):
            target = np.array([0.25 if step >= 200 else 0, 0, 0], dtype=np.float32)
            previous = act(env, policies, previous, target)
            env.step(previous)
            if (step + 1) % STRIDE == 0:
                env.render()
                frames = env.outputs[RGB("float32")].read()[0]
                views = panel(frames,
                              ["1 / Fly the preset - FPV", "Same flight - chase camera"])
                sensors = sensor_panel(env.outputs[Depth()].read()[0, 0], env.outputs[Segmentation(labels="semantic")].read()[0, 0],
                                       env.outputs[Normals()].read()[0, 0], env.outputs[Flow()].read()[0, 0], env.outputs[FlowStatus()].read()[0, 0])
                video.append_data(np.concatenate([views, sensors], axis=0))

show_video(output / "01-preset.mp4")

Rendering backend: optix
Visible semantic IDs: {96: 'FloorLamp', 127: 'Dresser', 159: 'Floor', 162: 'Wall', 4294967293: 'Agent'}

Multiple drone models

fpv_positions = {
    "x500": [0.1, 0, 0.19],
    "arpl": [0.08, 0, 0.08],
    "soft": [0.1, 0, 0.105],
    "crazyflie_brushless": [0.016, 0, 0.026],
    "crazyflie": [0.012, 0, 0.022],
    "savagebee_pusher": [0.035, 0, 0.035],
}
names = list(fpv_positions)
drones = [{"asset": rigged_model_path(name), "parameters": {
    "dynamics": {"dynamics": {"model": name}}, "cameras": {"front": camera(fpv_positions[name], 0.3, 0, 100)},
}} for name in names]
group_config = replace(config, num_agents=len(names) + 1, num_camera_agents=len(names) + 1)
with closing(HyperDrone(scenes, config=group_config, drones=drones,
                              parameters={"dynamics": initial})) as mixed:
    mixed.reset(spawn_points=shots["models"])
    mixed.set_instance_parameters({"agents": [
        {"cameras": {"front": camera([0.4, 0, 0.12], 0.35, 0, 90)}}, *[{} for _ in names],
    ]})
    anchors = np.array(shots["models"][0]["poses"], dtype=np.float32)[:, :3]
    policies, previous = pilots(mixed)
    with writer(output / "02-models.mp4", FPS) as video:
        for step in range(STEPS):
            targets = anchors.copy()
            targets[1:, 2] += 0.12 * np.sin(2 * np.pi * step * mixed.dt / 4 + np.arange(len(names)))
            previous = act(mixed, policies, previous, targets)
            mixed.step(previous)
            if (step + 1) % STRIDE == 0:
                filming_view = mixed.frames().reshape(1, mixed.num_views, HEIGHT, WIDTH, 3)[0, 0]
                video.append_data(panel([filming_view], ["Six drone models | filmed by a seventh drone"]))
show_video(output / "02-models.mp4")

Target-frame task

with closing(HyperDrone(scenes, config=replace(config, task="target_frame"),
                              drone_asset=drone_asset, parameters={"dynamics": initial, "cameras": {"front": fpv}})) as task:
    print("preset:", image_layout.blocks, "target_frame:", task.observation_layout.blocks)
    task.reset(spawn_points=shots["flight"])
    task.observe()
    target_offset, target_size = task.observation_layout.blocks["target_image"]
    target_image = task.frames()[0, ..., target_offset:target_offset + target_size].copy()
    policies, previous = pilots(task)
    with writer(output / "03-target.mp4", FPS) as video:
        for step in range(STEPS):
            previous = act(task, policies, previous)
            task.step(previous)
            if (step + 1) % STRIDE == 0:
                frame = task.frames()[0]
                images = [frame[..., offset:offset + size] for offset, size in task.observation_layout.blocks.values()]
                np.testing.assert_array_equal(images[1], target_image)
                video.append_data(panel(images, ["2 / Current image - returning to the goal", "Target image - fixed this episode"]))
show_video(output / "03-target.mp4")

preset: {'image': (0, 3)} target_frame: {'image_stack': (0, 3), 'target_image': (3, 3)}

Python task

with closing(HyperDrone(scenes, config=replace(config, instances_per_world=2), drone_asset=drone_asset,
                              parameters={"dynamics": initial, "cameras": {"front": fpv}})) as env:
    spawn_cycles = [[cycle(starts) for starts in room] for room in shots["python_task"]]
    room = 0
    env.reset(spawn_points=[next(starts) for starts in spawn_cycles[room]])
    policies, previous = pilots(env)
    targets = np.array([[0.2, 0, 0], [-0.2, 0, 0]], dtype=np.float32)
    age, episodes = np.zeros(2, dtype=int), np.zeros(2, dtype=int)
    last_reset = np.full(2, -100)
    with writer(output / "04-python-task.mp4", FPS) as video:
        for step in range(800):
            previous = act(env, policies, previous, targets)
            env.step(previous)
            state = env.observe_privileged()
            distance = np.linalg.norm(block(state, base_layout, "position") - targets, axis=1)
            speed = np.linalg.norm(block(state, base_layout, "linear_velocity"), axis=1)
            reward = -distance - 0.1 * speed
            age += 1
            terminated = ((distance < 0.04) & (speed < 0.15) & (age >= 100)) | (distance > 1.0)
            reset = terminated | (age >= [250, 350])
            if step == 399:
                env.rotate_scene()
                room = 1
                reset[:] = True
            env.reset(reset, spawn_points=[next(spawn_cycles[room][i]) for i in np.flatnonzero(reset)])
            for i in np.flatnonzero(reset):
                policies[i].reset()
                previous[i] = 0
                age[i] = 0
                episodes[i] += 1
                targets[i] *= -1
                last_reset[i] = step
            if (step + 1) % STRIDE == 0:
                env.observe()
                titles = [f"3 / Instance {i} | episode {episodes[i]} | " +
                          ("RESET" if step - last_reset[i] < 20 else f"reward {reward[i]:.2f} | error {distance[i]:.2f} m")
                          for i in range(2)]
                video.append_data(panel(env.frames(), titles))
    print("completed episodes:", episodes.tolist())
show_video(output / "04-python-task.mp4")

completed episodes: [8, 8]

Custom C++ world

%%writefile my_world.h
#pragma once
#include <rl_tools/rl/environments/hyperdrone/presets.h>
#include <cstddef>

namespace hyperdrone_user {
    struct SPEC: rl_tools::rl::environments::hyperdrone::presets::X500FPV<float, std::size_t> {
        static constexpr std::size_t N_CAMERA_AGENTS = 1;
        static constexpr std::size_t INSTANCES_PER_WORLD = 1;
        static constexpr std::size_t CAM_WIDTH = 640;
        static constexpr std::size_t CAM_HEIGHT = 360;
        using SHADING = rl_tools::hyperlux::VeryHigh;
        static constexpr bool ENABLE_MOTION_BLUR = true;
        static constexpr std::size_t MOTION_BLUR_SAMPLES = 8;
        static constexpr bool ENABLE_ANTI_ALIASING = true;
        static constexpr std::size_t ANTI_ALIASING_GRID_SIZE = 2;
    };
    using WORLD = rl_tools::rl::environments::hyperdrone::World<SPEC>;
}

with closing(HyperDrone(scenes, config=Config(spec_header="my_world.h"),
                              drone_asset=drone_asset, parameters={"dynamics": initial, "cameras": {"camera0": chase}})) as custom:
    custom.reset(spawn_points=shots["flight"])
    policies, previous = pilots(custom)
    with writer(output / "05-custom-world.mp4", FPS) as video:
        for step in range(STEPS):
            target = np.array([0.25 if step >= 200 else 0, 0, 0], dtype=np.float32)
            previous = act(custom, policies, previous, target, layout=base_layout)
            custom.step(previous)
            if (step + 1) % STRIDE == 0:
                views = custom.observe().reshape(custom.total_instances, HEIGHT, WIDTH, custom.image_channels)
                video.append_data(panel([views[0]], ["4 / Custom C++ world"]))

show_video(output / "05-custom-world.mp4")

Combined video

clips = [output / name for name in ("01-preset.mp4", "02-models.mp4", "03-target.mp4",
                                   "04-python-task.mp4", "05-custom-world.mp4")]
join(clips, "hyperdrone.mp4", FPS)
show_video("hyperdrone.mp4")
print("Saved", Path("hyperdrone.mp4").resolve())

Saved /home/jonas/rl-tools2/docs/hyperdrone/_build/tmp/tmpxdxtvoe5/hyperdrone.mp4