Platform for General Robot Intelligence Development
# baisc drone operations

import airgen

client = airgen.connect_airgen(robot_type="multirotor")
client.enableApiControl(True)

# Take off

client.takeoffAsync().join()

# Move to an altitude of 25 meters

client.moveToZAsync(-25, 5).join()

# Fly in a square path using velocity control

client.moveByVelocityAsync(5, 0, 0, 5).join()
client.moveByVelocityAsync(0, 5, 0, 5).join()
client.moveByVelocityAsync(-5, 0, 0, 5).join()
client.moveByVelocityAsync(0, -5, 0, 5).join()

client.hoverAsync().join()

# Fly in a square path using waypoints

waypoints = [
    airgen.Vector3r(25, 0, -25),
    airgen.Vector3r(25, 25, -25),
    airgen.Vector3r(0, 25, -25),
    airgen.Vector3r(0, 0, -25),
]

client.moveOnPathAsync(
    waypoints,
    5,
    60,
    airgen.DrivetrainType.MaxDegreeOfFreedom,
    airgen.YawMode(False, 0),
    -1,
    1,
).join()

# Fly in a square path with yaw angle relative to the direction of travel

client.moveOnPathAsync(
    waypoints, 5, 60, airgen.DrivetrainType.ForwardOnly, airgen.YawMode(False, 0), -1, 1
).join()

# Land

client.moveToZAsync(-2, 5).join()
client.landAsync().join()

airgen-samples/01_hello_drone.py
ScaledFoundations/GRID-playground
# reading images from camera in AirGen

from typing import List, Any, Literal, Tuple, Dict, Callable
import os
import math
import numpy as np
import rerun as rr

import airgen
from airgen.utils.sensor import imagetype2request, responses2images

client = airgen.connect_airgen(robot_type="multirotor")
client.enableApiControl(True)
client.takeoffAsync().join()
client.moveToPositionAsync(0, 0, -10, 5).join()
rr.init("airgen/camera")
rr.serve()

## part I: read images from camera in AirGen
print("Part I: read images from camera in AirGen\n")
# read every type of airgen image except for segmentation, which is discussed next
imagetypes = [
    airgen.ImageType.Scene,
    airgen.ImageType.DepthPlanar,
    airgen.ImageType.DepthPerspective,
    airgen.ImageType.DepthVis,
    airgen.ImageType.DisparityNormalized,
    airgen.ImageType.SurfaceNormals,
    airgen.ImageType.Infrared,
    airgen.ImageType.OpticalFlow,
    airgen.ImageType.OpticalFlowVis,
]
images = []
image_requests = [
    imagetype2request("front_center", image_type) for image_type in imagetypes
]

image_responses = client.simGetImages(requests=image_requests)
images = responses2images(image_responses)

for i, (image, _) in enumerate(images):
    print(
        f"{imagetypes[i]}'s image as numpy array: type={type(image)}, dtype={image.dtype}, shape={image.shape}"
    )


# visualize the images in rerun
# visualize the images
rr.log_image("airgen/camera/rgb", images[0][0])
rr.log_depth_image("airgen/camera/depth_perspective", images[1][0])
rr.log_depth_image("airgen/camera/depth_planar", images[2][0])
rr.log_image("airgen/camera/depth_vis", images[3][0])
rr.log_segmentation_image("airgen/camera/normalized_disparity", images[4][0])
rr.log_image(
    "airgen/camera/infrared", images[6][0]
)  # todo: infrared image is not working
rr.log_image("airgen/camera/optical_flow_vis", images[8][0])

## part II: read segmentation image from camera in AirGen
print("Part II: read segmentation image from camera in AirGen\n")
# Segmentation images capture the attribute: `SegmentationID` assigned to each object in the simulation scene.
# Initially, each object (mesh) is assigned random `SegmentationID` (`uint8`). When we request a
# `Segmentation` image from AirGen simulator, the camera captures the `SegmentationID` of all objects in the scene.
# One main usage of segmentation image is to extract specific objects
# or instances from the scene for processing or training of models.
# 1. the first step of using segmentation image is usually setting all objects' `SegmentationID` to 0,
# 2. then setting specific objects (through object name) to some non-zero values.
# 3. request the segmentation image from AirGen simulator.
#       The raw image data in response is represented as a numpy array with `dtype=uint8` and `shape=(height, width, 3)`,
#       where each RGB channel (or tuple of RGB values) corresponds to a fixed `SegmentationID` between `[0, 255]`.
# 4. finally we map the RGB values to `SegmentationID` to get the segmentation image.

# see all objects in the environment

objects = client.simListSceneObjects()
print("All objects in the environment:", objects)

# set all objects' `SegmentationID` to 0
client.simSetSegmentationObjectID("[\w]*", 0, is_name_regex=True)

# set the SegmentationID of all objects to 1
client.simSetSegmentationObjectID("[\w]*", 1, is_name_regex=True)


requests = [imagetype2request("front_center", airgen.ImageType.Segmentation)]
responses = client.simGetImages(requests=requests)
# in the helper function:  `responses2image`, the raw segmentation image (in RGB format) is converted to a numpy array of `SegmentationIDs`
segmentation_image = responses2images(responses)[0][0]
print(
    f"Segmentation image as numpy array: type={type(segmentation_image)}, dtype={segmentation_image.dtype}, shape={segmentation_image.shape}"
)
rr.log_segmentation_image("airgen/camera/segmentation", segmentation_image)

## part III: visualize trajectory along with camera captures in 3D space
print("Part III: visualize trajectory along with camera captures in 3D space\n")

from airgen.utils.mechanics import depth2pointcloud
from airgen.utils.collect import data_collector
from airgen.utils.visualize import rr_camera_capture_logger


def collect_camera_capture(client) -> dict:
    # use rr.log with `time_less=True` to log data inside data collection task
    requests = [
        imagetype2request("front_center", airgen.ImageType.DepthPerspective),
        imagetype2request("front_center", airgen.ImageType.Scene),
        imagetype2request("bottom_center", airgen.ImageType.DepthPerspective),
        imagetype2request("bottom_center", airgen.ImageType.Scene),
    ]
    responses = client.simGetImages(requests=requests)
    images = responses2images(responses)
    ## can also visualize inside rerun, let's leave that to the user
    return {
        "front_depth": images[0],
        "front_rgb": images[1],
        "bottom_depth": images[2],
        "bottom_rgb": images[3],
    }


@data_collector(collect_camera_capture, time_delta=0.2)
def move_task(client, **kwargs) -> None | Tuple[None, list]:
    client.moveToPositionAsync(0, 0, -10, 5).join()
    client.moveToPositionAsync(10, 10, -10, 5).join()


_, data = move_task(client, _collect_data=True)


capture_logger = rr_camera_capture_logger("airgen/world")
for i, capture in enumerate(data):
    for camera_name, camera_data in capture.items():
        if i == 0 and camera_name.endswith("depth"):
            # compute point cloud and visualize it in rerun
            # remove points that are too far away
            mask = np.where(camera_data[0] < 1000.0, 1, 0).astype(np.uint8)
            point_cloud = depth2pointcloud(
                depth=camera_data[0], camera_param=camera_data[1], mask=mask
            )
            rr.log_points(f"airgen/world/point_cloud/{camera_name}", point_cloud)
        # camera_data is a tuple of (image, camera_params)
        if camera_name.endswith("rgb"):
            capture_logger(
                camera_data[0], camera_data[1], camera_name=f"{camera_name}_{i}"
            )
# shutdown rerun server
rr.disconnect()
rr.rerun_shutdown()

airgen-samples/02_camera.py
ScaledFoundations/GRID-playground
# read sensor data in AirGen

from typing import List, Tuple, Dict, Any, Optional, Callable
import numpy as np
import rerun as rr

import airgen

from airgen.utils.collect import data_collector

# setup client
client = airgen.connect_airgen(robot_type="multirotor")
client.enableApiControl(True)
client.takeoffAsync().join()

rr.init("airgen/sensor")
rr.serve()


def readSensors(client: airgen.VehicleClient) -> dict:
    sensor_data = {}
    sensor_data["barometer"] = client.getBarometerData()
    sensor_data["imu"] = client.getImuData()
    sensor_data["gps"] = client.getGpsData()

    sensor_data["magnetometer"] = client.getMagnetometerData()
    sensor_data["distance1"] = client.getDistanceSensorData(
        distance_sensor_name="Distance1"
    )
    sensor_data["distance2"] = client.getDistanceSensorData(
        distance_sensor_name="Distance2"
    )
    sensor_data["lidar"] = client.getLidarData()
    return sensor_data


@data_collector(readSensors, time_delta=4.0)
def move_task(
    client: airgen.MultirotorClient, position: Tuple[float], **kwargs
) -> None | Tuple[None, List[dict]]:
    client.moveToPositionAsync(*position, 1).join()


_, sensor_data = move_task(client, (0, 0, -10), _collect_data=True)
for i, data in enumerate(sensor_data):
    lidar = data["lidar"]
    rr.log_points(
        f"airgen/sensor/lidar_{i}",
        np.array(lidar.point_cloud).reshape(-1, 3),
        class_ids=lidar.segmentation,
    )
print(f"collected {len(sensor_data)} measurements during moving task")

rr.disconnect()
rr.rerun_shutdown()

airgen-samples/03_sensor.py
ScaledFoundations/GRID-playground
import airgen
import numpy as np


class DroneController:
    def __init__(self):
        # Initialize a drone client
        self.drone_client = airgen.connect_airgen(robot_type="multirotor")

        # Reset previous state if any
        self.drone_client.reset()

        # Confirm the connection, and enable offboard (API) control
        self.drone_client.confirmConnection()
        self.drone_client.enableApiControl(True)

        # Take off
        self.drone_client.takeoffAsync().join()

        # Get navigation map information from the simulation world
        self.nav_mesh_info = self.drone_client.getNavMeshInfo()
        print("Nav mesh info: {}".format(self.nav_mesh_info))

    def sample_random_pose(self):
        # calculate bounds of the nav mesh
        amplitude = (
            np.absolute(
                np.array(
                    [
                        self.nav_mesh_info[2]["x_val"] - self.nav_mesh_info[1]["x_val"],
                        self.nav_mesh_info[2]["y_val"] - self.nav_mesh_info[1]["y_val"],
                        self.nav_mesh_info[2]["z_val"] - self.nav_mesh_info[1]["z_val"],
                    ]
                )
            )
            / 2.0
        )
        # sample a random point on the nav mesh
        random_point = airgen.Vector3r(
            np.random.uniform(
                self.nav_mesh_info[0]["x_val"] - amplitude[0],
                self.nav_mesh_info[0]["x_val"] + amplitude[0],
            ),
            np.random.uniform(
                self.nav_mesh_info[0]["y_val"] - amplitude[1],
                self.nav_mesh_info[0]["y_val"] + amplitude[1],
            ),
            np.random.uniform(
                self.nav_mesh_info[0]["z_val"] - amplitude[2],
                self.nav_mesh_info[0]["z_val"] + amplitude[2],
            ),
        )
        # sample a random yaw angle
        random_yaw = np.random.uniform(-np.pi, np.pi)
        return airgen.Pose(
            random_point, airgen.Quaternionr(airgen.Vector3r(0, 0, random_yaw))
        )

    def plan_and_move(self):
        # Get current pose of the drone in 6-DOF
        start_pose = self.drone_client.simGetVehiclePose()

        # Sample a random valid pose in the environment
        goal_pose = self.sample_random_pose()
        goal_pose.position.z_val = start_pose.position.z_val

        # Compute a collision-free path between start and goal
        trajectory = self.drone_client.simPlanPath(
            start_pose.position, goal_pose.position, True, True
        )

        # always check the lenght of the trajectory before moving
        # if len(trajectory) < 2, then no viable path was found
        # in this case, consider using:
        #   drone_client.simPlanPathToRandomizedGoal: find a path with end point within a certain radius of the goal
        #   drone_client.simPlanPathToRandomFreePoint: find a path with end point within a certain radius of a random point (current position for example)

        points = []
        for waypoint in trajectory:
            points.append(
                airgen.Vector3r(waypoint["x_val"], waypoint["y_val"], waypoint["z_val"])
            )

        # Move the drone along the planned path at a velocity of 5 m/s
        velocity = 5.0
        self.drone_client.moveOnPathAsync(
            points,
            velocity,
            120,
            airgen.DrivetrainType.ForwardOnly,
            airgen.YawMode(False, 0),
            -1,
            0,
        ).join()


controller = DroneController()

while True:
    try:
        pose = controller.sample_random_pose()
        print(
            f"Moving to new position: [{pose.position.x_val}, {pose.position.y_val}, {pose.position.z_val}]"
        )
        controller.plan_and_move()
    except KeyboardInterrupt:
        print("Keyboard interrupt detected, exiting...")
        client = airgen.MultirotorClient()
        client.reset()
        break

airgen-samples/04_path_planning.py
ScaledFoundations/GRID-playground
import airgen
import threading


def run_drone(id: int):
    client = airgen.connect_airgen(robot_type="multirotor")
    client.confirmConnection()

    vehicle_name = f"Drone_{id}"
    pose = airgen.Pose(airgen.Vector3r(0, 5 * id, 0), airgen.Quaternionr(0, 0, 0, 0))

    print(f"Creating {vehicle_name}")
    success = client.simAddVehicle(vehicle_name, "SimpleFlight", pose)
    client.enableApiControl(True, vehicle_name)

    if not success:
        print(f"Failed to create {vehicle_name}")
        return

    client.moveToZAsync(-25, 5, vehicle_name=vehicle_name)


if __name__ == "__main__":
    num_vehicles = 5

    print(f"Creating {num_vehicles} vehicles")

    threads = []
    for id in range(num_vehicles, 0, -1):
        t = threading.Thread(target=run_drone, args=(id,))
        threads.append(t)
        t.start()

    for t in threads:
        t.join()

airgen-samples/05_multi_drone.py
ScaledFoundations/GRID-playground
import airgen
import time
client = airgen.connect_airgen(robot_type="multirotor")

client.moveByVelocityAsync(0, -5, 0, 120)
client.simEnableWeather(False)

# Iterate over hour of day from 12 to 23 and then 0 to 12

for hour in range(12, 24):
    for minute in range(0, 60):
        print(f"2023-09-08 {hour}:{minute}:00")
        client.simSetTimeOfDay(True, f"2023-09-08 {hour}:{str(minute).zfill(2)}:00", True)
        time.sleep(0.016)
for hour in range(0, 12):
    for minute in range(0, 60):
        print(f"2023-09-08 {hour}:{minute}:00")
        client.simSetTimeOfDay(True, f"2023-09-08 {hour}:{str(minute).zfill(2)}:00", True)
        time.sleep(0.016)
    
client.simEnableWeather(True)

for i in range(0, 10):
    client.simSetWeatherParameter(airgen.WeatherParameter.Rain, i/10)
    time.sleep(0.1)
for i in range(10, 0, -1):
    client.simSetWeatherParameter(airgen.WeatherParameter.Rain, i/10)
    time.sleep(0.1)
for i in range(0, 10):
    client.simSetWeatherParameter(airgen.WeatherParameter.Snow, i/10)
    client.simSetWeatherParameter(airgen.WeatherParameter.Fog, i/10)
    time.sleep(0.1)
for i in range(10, 0, -1):
    client.simSetWeatherParameter(airgen.WeatherParameter.Snow, i/10)
    client.simSetWeatherParameter(airgen.WeatherParameter.Fog, i/10)
    time.sleep(0.1)
for i in range(0, 10):
    client.simSetWeatherParameter(airgen.WeatherParameter.Dust, i/10)
    time.sleep(0.1)
for i in range(10, 0, -1):
    client.simSetWeatherParameter(airgen.WeatherParameter.Dust, i/10)
    time.sleep(0.1)


client.simSetWind(airgen.Vector3r(5.0, 0, 0))
time.sleep(5)
client.simSetWind(airgen.Vector3r(0, 0, 0))
airgen-samples/06_weather.py
ScaledFoundations/GRID-playground
# this script is to be run with AirGen binary built with google maps support
# before starting the binary, make sure you `settings.json` has the following entry (in the top level): "GMapsApiKey": "your_google_map_api_key"
import airgen

# Note: To use `airgen.connect_airgen(robot_type="multirotor", geo=True)`, make sure the latest airgen (>=0.0.6) is installed by (pip install airgen --upgrade)
client = airgen.connect_airgen(robot_type="multirotor", geo=True)
# otherwise, use the following two lines instead (and comment out the line above)
# client = airgen.MultirotorClient(geo=True)
# client.confirmConnection()
client.enableApiControl(True)

# Teleport vehicle to the vicinity of Seattle
geopose = airgen.GeoPose(
    airgen.GeoPoint(47.62542093668409, -122.36989783619246, 200), airgen.Quaternionr()
)
client.simSetVehicleGeoPose(geopose, True)

# Note: If you teleport the agent too far from the OriginGeoPoint mentioned in settings.json,
# it may not work properly. In that case, you should change the reference point of the world to a closer
# lat-long coordinates. You can do that either in OriginGeoPoint, or at run time through
# client.simSetGeoReference(airgen.GeoPoint(lat, long, alt))

# Takeoff
client.takeoffAsync().join()

print(client.getGpsData())

# Fly to a nearby destination
goal = airgen.GeoPoint(47.61168088062645, -122.32718628892, 200)
client.moveToGPSAsync(
    goal,
    25.0,
    timeout_sec=3e38,
    drivetrain=airgen.DrivetrainType.ForwardOnly,
    yaw_mode=airgen.YawMode(False, 0),
).join()

print(client.getGpsData())

airgen-samples/07_geospecific.py
ScaledFoundations/GRID-playground
# collision
# collision points are designed for Block environment

from typing import Tuple, List
import time
import airgen

# import a function decorator that monitors collision
# see documentation for more details and usage
from airgen.utils.collect import collision_collector


client = airgen.connect_airgen(robot_type="multirotor")
client.enableApiControl(True)
client.takeoffAsync().join()


@collision_collector()
def move(client, x, y, z, **kwargs) -> None | Tuple[None, list]:
    future = client.moveToPositionAsync(x, y, z, 3)
    future.start_waiting()
    while future.is_completed is None:
        # waiting it to complete
        time.sleep(1)


@collision_collector()
def move2(client, x, y, z, **kwargs) -> None | Tuple[None, list]:
    # or just use below one-liner, which is equivalent to the above move
    client.moveToPositionAsync(x, y, z, 3).join()


# todo: fix the linter error
_, collisions = move(
    client, 0, 0, -10, _stop_at_collision=True, _collect_collision=True
)  # noqa
print(collisions)

_, collisions = move(
    client,
    20,
    20,
    -10,
    _stop_at_collision=True,
    _collect_collision=True,
)
print(collisions)

# Note
#   if no explict action is applied to drone after collision (like `__stop_at_collision__`, which puts drone in hover state),
#   The behavior of drone is not clearly defined, it will try to recover from collision, but it also start drifting to a random direction.

airgen-samples/08_watch_collision.py
ScaledFoundations/GRID-playground
# Visual Odometry using the AirGen simulator. 
# This script expects the AirGen simulator to be set to "ComputerVision" mode.
# Uses DPVO to compute up-to-scale poses from RGB data.

import numpy as np
import rerun as rr
import airgen
from airgen.utils.sensor import imagetype2request, responses2images
from grid import GRIDConfig
from grid.model.perception.vo.vo_dpvo import VO

# Point this to the directory that contains `external`
GRIDConfig.set_main_dir("/GRID")

rr.init('vo')
rr.serve()

# Assuming a 640x480 camera with 90 degree FoV. This should match the
# setting in the AirGen settings.json
calib = np.array([320, 320, 320, 240])
dpvo = VO(calib)

c = airgen.connect_airgen(robot_type="vehicle")
c.confirmConnection()

t = 0

imagetypes = [airgen.ImageType.Scene]
image_requests = [
    imagetype2request("front_center", image_type) for image_type in imagetypes
]

while True:
    try:
        responses = c.simGetImages(requests = image_requests)
        image = responses2images(responses)[0][0]
        current_pose = dpvo.run(image)
        rr.log_point(f'vo/position/pos_{t}', [current_pose[0], current_pose[1], -current_pose[2]])

        t += 1
    except KeyboardInterrupt:
        pred_traj = dpvo.terminate()
        break

grid-samples/visual_odometry.py
ScaledFoundations/GRID-playground