UAV Models and Usage#

See also

We strongly recommend going through the tutorials provided by Isaac Sim and Isaac Orbit.

Using the UAV Models#

All UAV models, e.g., Hummingbird and Firefly are subclassed from MultirotorBase. An instance of MultirotorBase does not correspond to a specifc UAV, but acts like a view (think of torch.Tensor.view) that holds the states of a group of UAVs of that type in torch.Tensor for reading and writing. It implements the multirotr dynamics and provides interfaces for common interactions such as getting the kinematic information and applying rotor commands.

The following example demonstrates:

  1. How to create and control a UAV using the MultirotorBase class.

  2. How to use a position controller to track a circular trajectory.

  3. How to attach cameras to UAVs to record images and hence videos.

examples/demo_drone.py#
  1import torch
  2
  3import hydra
  4from omegaconf import OmegaConf
  5from omni_drones import CONFIG_PATH, init_simulation_app
  6
  7
  8@hydra.main(version_base=None, config_path=".", config_name="demo")
  9def main(cfg):
 10    OmegaConf.resolve(cfg)
 11    simulation_app = init_simulation_app(cfg)
 12    print(OmegaConf.to_yaml(cfg))
 13
 14    import omni_drones.utils.scene as scene_utils
 15    from omni.isaac.core.simulation_context import SimulationContext
 16    from omni_drones.controllers import LeePositionController
 17    from omni_drones.robots.drone import MultirotorBase
 18    from omni_drones.utils.torch import euler_to_quaternion
 19    from omni_drones.sensors.camera import Camera, PinholeCameraCfg
 20    import dataclasses
 21
 22    sim = SimulationContext(
 23        stage_units_in_meters=1.0,
 24        physics_dt=cfg.sim.dt,
 25        rendering_dt=cfg.sim.dt,
 26        sim_params=cfg.sim,
 27        backend="torch",
 28        device=cfg.sim.device,
 29    )
 30    n = 4
 31
 32    drone_cls = MultirotorBase.REGISTRY[cfg.drone_model]
 33    drone = drone_cls()
 34
 35    translations = torch.zeros(n, 3)
 36    translations[:, 1] = torch.arange(n)
 37    translations[:, 2] = 0.5
 38    drone.spawn(translations=translations)
 39
 40    scene_utils.design_scene()
 41
 42    camera_cfg = PinholeCameraCfg(
 43        sensor_tick=0,
 44        resolution=(320, 240),
 45        data_types=["rgb"],
 46    )
 47    # cameras used as sensors
 48    camera_sensor = Camera(camera_cfg)
 49    camera_sensor.spawn([
 50        f"/World/envs/env_0/{drone.name}_{i}/base_link/Camera" 
 51        for i in range(n)
 52    ])
 53    # camera for visualization
 54    camera_vis = Camera(dataclasses.replace(camera_cfg, resolution=(960, 720)))
 55
 56    sim.reset()
 57    camera_sensor.initialize(f"/World/envs/env_0/{drone.name}_*/base_link/Camera")
 58    camera_vis.initialize("/OmniverseKit_Persp")
 59    drone.initialize()
 60
 61    # let's fly a circular trajectory
 62    radius = 1.5
 63    omega = 1.
 64    phase = torch.linspace(0, 2, n+1, device=sim.device)[:n]
 65
 66    def ref(t):
 67        _t = phase * torch.pi + t * omega
 68        pos = torch.stack([
 69            torch.cos(_t) * radius,
 70            torch.sin(_t) * radius,
 71            torch.ones(n, device=sim.device) * 1.5
 72        ], dim=-1)
 73        vel_xy = torch.stack([
 74            -torch.sin(_t) * radius * omega,
 75            torch.cos(_t) * radius * omega,
 76        ], dim=-1)
 77        yaw = torch.atan2(vel_xy[:, 1], vel_xy[:, 0])
 78        return pos, yaw
 79
 80    init_rpy = torch.zeros(n, 3, device=sim.device)
 81    init_pos, init_rpy[:, 2] = ref(torch.tensor(0.0).to(sim.device))
 82    init_rot = euler_to_quaternion(init_rpy)
 83    init_vels = torch.zeros(n, 6, device=sim.device)
 84
 85    # create a position controller
 86    # note: the controller is state-less (but holds its parameters)
 87    controller = LeePositionController(g=9.81, uav_params=drone.params).to(sim.device)
 88
 89    def reset():
 90        drone._reset_idx(torch.tensor([0]))
 91        drone.set_world_poses(init_pos, init_rot)
 92        drone.set_velocities(init_vels)
 93        # flush the buffer so that the next getter invocation 
 94        # returns up-to-date values
 95        sim._physics_sim_view.flush() 
 96    
 97    reset()
 98    drone_state = drone.get_state()[..., :13].squeeze(0)
 99
100    frames_sensor = []
101    frames_vis = []
102    from tqdm import tqdm
103    for i in tqdm(range(1000)):
104        if sim.is_stopped():
105            break
106        if not sim.is_playing():
107            sim.render()
108            continue
109        ref_pos, ref_yaw = ref((i % 1000)*cfg.sim.dt)
110        action = controller(drone_state, target_pos=ref_pos, target_yaw=ref_yaw)
111        drone.apply_action(action)
112        sim.step(render=True)
113
114        if i % 2 ==  0:
115            frames_sensor.append(camera_sensor.get_images().cpu())
116            frames_vis.append(camera_vis.get_images().cpu())
117
118        if i % 1000 == 0:
119            reset()
120        drone_state = drone.get_state()[..., :13].squeeze(0)
121
122    from torchvision.io import write_video
123
124    for image_type, arrays in torch.stack(frames_sensor).items():
125        print(f"Writing {image_type} of shape {arrays.shape}.")
126        for drone_id, arrays_drone in enumerate(arrays.unbind(1)):
127            if image_type == "rgb":
128                arrays_drone = arrays_drone.permute(0, 2, 3, 1)[..., :3]
129                write_video(f"demo_rgb_{drone_id}.mp4", arrays_drone, fps=1/cfg.sim.dt)
130            elif image_type == "distance_to_camera":
131                arrays_drone = -torch.nan_to_num(arrays_drone, 0).permute(0, 2, 3, 1)
132                arrays_drone = arrays_drone.expand(*arrays_drone.shape[:-1], 3)
133                write_video(f"demo_depth_{drone_id}.mp4", arrays_drone, fps=0.5/cfg.sim.dt)
134
135    for image_type, arrays in torch.stack(frames_vis).items():
136        print(f"Writing {image_type} of shape {arrays.shape}.")
137        for _, arrays_drone in enumerate(arrays.unbind(1)):
138            if image_type == "rgb":
139                arrays_drone = arrays_drone.permute(0, 2, 3, 1)[..., :3]
140                write_video(f"demo_rgb.mp4", arrays_drone, fps=1/cfg.sim.dt)
141            elif image_type == "distance_to_camera":
142                arrays_drone = -torch.nan_to_num(arrays_drone, 0).permute(0, 2, 3, 1)
143                arrays_drone = arrays_drone.expand(*arrays_drone.shape[:-1], 3)
144                write_video(f"demo_depth.mp4", arrays_drone, fps=0.5/cfg.sim.dt)
145
146    simulation_app.close()
147
148
149if __name__ == "__main__":
150    main()

Importing/Creating New UAVs#

Since MultirotorBase implements the commonly use multirotor dynamics, adding an ordinary UAV is easy:

from omni_drones.robots.drone.multirotor import MultirotorBase
from omni_drones.robots.robot import ASSET_PATH


class Iris(MultirotorBase):

    usd_path: str = ASSET_PATH + "/usd/iris.usd"
    param_path: str = ASSET_PATH + "/usd/iris.yaml"

which effectively comes to 1. providing a description file in .usd format, and 2. specifying its parameters.

Adding an UAV model with more complex dynamics may require extending the corresponding methods. For example, the omnidirectional UAV Omav has 6 tilt units. So we extend Omav.apply_action() to control the tilt units via velocity targets, and Omav._reset_idx() to reset the tilting angle.

 1import torch
 2from functorch import vmap
 3from torchrl.data import BoundedTensorSpec, UnboundedContinuousTensorSpec
 4
 5from omni_drones.robots.drone import MultirotorBase
 6from omni_drones.robots.robot import ASSET_PATH
 7
 8class Omav(MultirotorBase):
 9
10    usd_path: str = ASSET_PATH + "/usd/omav.usd"
11    param_path: str = ASSET_PATH + "/usd/omav.yaml"
12
13    def __init__(self, name: str = "Omav", cfg=None) -> None:
14        super().__init__(name, cfg)
15        self.action_spec = BoundedTensorSpec(-1, 1, 12 + 6, device=self.device)
16        self.tilt_dof_indices = torch.arange(0, 6, device=self.device)
17        self.rotor_dof_indices = torch.arange(6, 18, device=self.device)
18        self.max_tilt_velocity = 10
19
20        self.action_spec = BoundedTensorSpec(-1, 1, self.num_rotors + 6, device=self.device)
21        self.state_spec = UnboundedContinuousTensorSpec(
22            19 + self.num_rotors + 12, device=self.device
23        )
24
25    def initialize(self, prim_paths_expr: str = None):
26        if not self.is_articulation:
27            raise NotImplementedError
28        super().initialize(prim_paths_expr)
29        self.init_joint_positions = self._view.get_joint_positions()
30        self.init_joint_velocities = self._view.get_joint_velocities()
31    
32    def apply_action(self, actions: torch.Tensor) -> torch.Tensor:
33        rotor_cmds, tilt_cmds = actions.expand(*self.shape, 18).split([12, 6], dim=-1)
34        super().apply_action(rotor_cmds)
35
36        velocity_targets = tilt_cmds.clamp(-1, 1) * self.max_tilt_velocity
37        self._view.set_joint_velocity_targets(
38            velocity_targets, joint_indices=self.tilt_dof_indices
39        )
40
41        return self.throttle.sum(-1)
42
43    def _reset_idx(self, env_ids: torch.Tensor, train: bool=True):
44        env_ids = super()._reset_idx(env_ids, train)
45        self._view.set_joint_positions(
46            self.init_joint_positions[env_ids],
47            env_indices=env_ids,
48        )
49        self._view.set_joint_velocities(
50            self.init_joint_velocities[env_ids],
51            env_indices=env_ids,
52        )
53        return env_ids
54
55    def get_state(self, env=True):
56        state = super().get_state(env)
57        joint_positions = self.get_joint_positions()[..., self.tilt_dof_indices]
58        joint_velocities = self.get_joint_velocities()[..., self.tilt_dof_indices]
59        state = torch.cat([state, joint_positions, joint_velocities], dim=-1)
60        assert not torch.isnan(state).any()
61        return state
62

Customized Dynamics#

Note

To be refactored.

Building New Configurations#

Note

To be refactored.