Imitation Learning Policy

Imitation learning is a type of machine learning where a model is trained to imitate the behavior of an expert. In this case, the expert is the driver of the car. The model is trained to imitate the behavior of the driver by learning from the dataset collected from manual driving.

This tutorial will guide you through the process of training an imitation learning policy using a dataset collected from manual driving. The main steps are:

1. Collect a dataset of expert driving demonstrations suitable for training a policy.
2. Process the data to create individual training examples.
3. Train a neural network to imitate the behavior of the expert.
4. Deploy the trained network on the autonomous vehicle and evaluate the policy's behavior.

Prerequisites

• Ensure that you've calibrated your car and gotten the Navigation Stack working, including building the workspace on the car with scripts/build. Data collection and rollout are airfield plans, and their panes launch from that shared build. Plain scripts/build only compiles the navstack packages, so also build the ones this tutorial adds: scripts/build av_recorder leg_detector av_imitation. (A pane compiles a missing package the first time it starts, so skipping this still works, just with a slow first launch.)
• The av_imitation image must include the GPU ML stack (torch, onnxruntime, tensorboard). On a migrated car, airfield package build av_imitation installs it; confirm with airfield package shell av_imitation and python3 -c "import torch; print(torch.cuda.is_available())", which should print True.

Airfield requires a workspace that has been migrated to it. Check for a plans/ folder and the scripts/build file in ~/roboracer_ws. If your car has a tmux/ folder and a container script instead, it is still on the older layout and you should follow the fallback instructions below. The full airfield reference for this workspace is the Airfield guide.

Learning Objectives

• Students can collect a dataset of expert driving demonstrations suitable for training a policy.
• Understand data processing to transform the recorded data into data suitable for training a policy.
• Gain hands-on experience with the neural network training process.
• Learn how to roll out the trained policy on the autonomous vehicle and evaluate the policy's behavior.

Tutorial Steps

We will use behavior cloning to train an imitation learning policy.

In particular, we are going to train a neural network to approximate the behavior of the driver. The neural network will take as input a short sequence of camera images and output vehicle velocity and steering angle commands.

We will record data by having an expert demonstrator (you) drive the car around while recording images and joystick commands.

The process of behavior cloning will teach the neural network to imitate your commands given the images as input. For example, imagine as you drive the car around that all you can see is the image from the camera. How would you control the car? This is what we will be teaching the neural network.

This process requires several careful considerations. First, the car will only be able to approximate the behavior that it has seen. That is, if it sees a new image, unlike anything it's seen before, it won't know what to do and could output a command that seems random. This is the problem of "out of distribution" data. The new image is outside of the distribution of images that the policy has seen during the training process. Second, how well you control the car will determine how well the policy can drive. For example, if you are training a policy that is intended to drive straight down a sidewalk, but you weave back and forth across the edges of the sidewalk during your recording, the car will learn to weave back and forth across the edges of the sidewalk as well.

1. Data Collection

The process used to collect your dataset and the specific properties of the data will make a large difference in the performance of your policy. Therefore, you should read through this tutorial before beginning the data collection process. The data you collect will define the cars behavior.

The nodes for recording are configured in the record_sensors plan: the recorder, the camera, the leg detector, the lidar, the vesc driver, the joystick, the robot description, the on-car GUI, and the VNC desktop. Launch it from the GUI on the car or ssh into your car and run:

cd ~/roboracer_ws
airfield project up record_sensors

There is no container shell to enter first. Each pane starts its own container from the package image it names, and the recorder, camera, and driver nodes are already installed in them.

The GUI will show "REC" followed by a red square in the center of the screen when the recorder is running, but not saving data. To toggle recording press the red circle on your joystick.

Drive the car around the track. The rosbags are written to ~/roboracer_ws/data/rosbags/ on the car, which persists across container restarts.

When you are done recording, tear the session down:

airfield project down

Always shut down with airfield project down rather than closing the tmux session, so the containers stop with it and the camera, lidar, and VESC serial port are released. After a hard crash or power loss, add --prune to sweep any leftover containers.

An alternative: record with tmuxinator

The car previously started the recorder from a tmuxinator configuration inside the container shell. This flow still works if airfield is unavailable, and it is the only flow available on a car that has not been migrated. The recording indicator and joystick toggle are unchanged.

cd ~/roboracer_ws
./container shell
cd ~/roboracer_ws/tmux/record_sensors/
tmuxinator

2. Data Processing

When we recorded our data, it was saved continuously, as it was published by the sensor nodes and received by the recorder. We now need to process this continuous data and transform it into individual training examples that can be used during neural network training. We will also apply data cleaning during this step to remove any data that is not relevant or is incorrect.

There are many choices to be made when transforming the recorded data into training examples, including:

• The data sampling rate
• How data is combined
• What data is extracted for

Our data was recorded as rosbag files. We will use the AV Imitation Annotator to transform the rosbag files into individual training examples.

While the commands here can be run on the car, it is much better to copy the data to a more powerful machine and run the data processing there.

In order to get the data from the car onto your computer, use the rsync command.

For example, on your development machine (not the car), run:

rsync -avP orin@orinXX:roboracer_ws/data/rosbags/ ~/roboracer_ws/data/rosbags/

Note that the orinXX hostname will only work on machines that have been properly configured and are connected to the VPN. Alternatively, use the WiFi IP address of the car instead of this hostname, found on the car's display or by running ip a on the car.

Once you've copied the data from your car to your development machine using the rsync command, start the AV Imitation Annotator webapp on your development machine. This tool will allow you to process the rosbag files, extract the relevant data and create a dataset suitable for training a policy.

The webapp runs in a container that only sees ~/roboracer_ws/data if your development machine's .air file shares that folder in. The setup script writes .air on a car, but on any other machine you create it yourself: see The .air file, and keep its ~/roboracer_ws/data line. Without it the Bags tab is empty.

On your development machine, start the webapp through airfield:

cd ~/roboracer_ws
airfield package cmd av_imitation -- ros2 run av_imitation start_webapp

Then in a web browser on your computer, go to the URL that start_webapp prints in the terminal (http://<machine>:<port>), as shown in the image.

An alternative: start the webapp with the container shell

On a car or dev machine that has not been migrated to airfield, start the webapp from inside the container shell instead:

cd ~/roboracer_ws
./container shell
ros2 run av_imitation start_webapp

The webapp has 5 tabs: Bags, Dataset, Visualize, Analysis, and Inference. Data processing uses the first 3; the Inference tab is covered in step 5.

1. In the Bags tab, you will see a list of rosbag files that were recorded on the car.

   

   Clicking on a rosbag file will allow you to preview the recorded data.

   Playing back the bag will allow you to view the recorded data and see what parts will be included in the dataset. The sections of the timeline marked in red are times when the L1 trigger was not depressed, and therefore will not be included in the dataset.

   You can enter any comments in the metadata texbox.

   In the Process Bag section, select the image processing options and then click "Start Processing" to extract the images and the actions from the rosbag file and save them as individual images in the dataset directory.

2. The Dataset tab will allow you to compile the extraced data into a dataset suitable for training a policy. Choose the sampling rates and duration of each example before and after the current timestep (you can use the default sampling rates). Then click "Generate Dataset" to create the dataset. This will create a metadata file which references the images and actions in the dataset directory.

   

3. The Visualize tab will allow you to visualize the dataset and check for any issues. You can select a dataset and click "Visualize" to view the dataset.

   

Note: Step 2 in the image above, the training command, uses --batch-size 32, which needs more memory than the Jetson Orin Nano has. Run that command as shown only on the lab computer. To train on the car, use the smaller batch sizes in the next section.

3. Neural Network Training

Note: Neural network training can take some time. Instead of trying to train the network with a single battery, plug in your Jetson Orin into the wall via a power adapter so it can stay on while training.

Training runs inside the av_imitation container. Open a shell in it:

cd ~/roboracer_ws
airfield package shell av_imitation

The shell starts in the package's source directory, so the training script is at av_imitation/src/main.py. First train the policy for 1 epoch with a batch size of 4, to test if the training process works:

Note: Replace outdoor_v1 with your dataset name. If --model mlp does not work try --model cnn

python3 av_imitation/src/main.py --model mlp --dataset outdoor_v1 --epochs 1 --batch-size 4
Starting experiment: 20251124_145208_mlp_outdoor_v1
Using device: cuda
Loading dataset metadata...
...
Output Steps: 2
...
Epoch [1/1] Train Loss: 0.0120 Val Loss: 0.0182
Saved best model to /home/orin/roboracer_ws/data/experiments/20251124_145208_mlp_outdoor_v1/models/best_model.pth
Training finished.

Using device: cuda confirms the GPU torch build from the prerequisites is in use. If it says cpu, the image has the CPU-only torch from airfield's shared recipe book instead of the project's GPU recipe, roboracer_ws/dependencies/arm64/torch.yaml, which takes priority over the shared one. Pull the latest roboracer_ws and av_imitation, then rebuild with airfield package build av_imitation. Run airfield status in packages/av_imitation to see which recipe each dependency comes from; see the Airfield dependencies section for how recipes are found.

Once the test training run successfully completes, you can train the policy for a longer duration by setting the --epochs flag to a higher value and the --batch-size flag to a higher value.

You'll want to pick a number of epochs greater than you will reach before early stopping. Set the --batch-size to a value that will allow the policy training to fit on your GPU, utilizing as much memory as possible, but not exceeding the available memory.

For example, to train the model:

Note: Training will use almost all the memory available on the Orin, so be sure to close any unnecessary processes. Use htop to check what is running and kill [pid] to stop any unnecessary processes. You can also reduce the batch size to fit on your GPU, however, this will slow down training and possibly reduce the performance of your model.

• Start a new tmux session outside of any container

• Split the window into 2 vertical panes

• In the left pane, open a container shell and run the training:

Note: Again, replace outdoor_v1 with your dataset name.

airfield package shell av_imitation
python3 av_imitation/src/main.py --model cnn --dataset outdoor_v1 --epochs 1000 --batch-size 4

• Split the right pane and in the top window run jtop (on the host, not inside a container):

jtop

• In the bottom window open another container shell and start tensorboard:

airfield package shell av_imitation
tensorboard --logdir ~/roboracer_ws/data/experiments --port 6006 --host 0.0.0.0

Note: The tensorboard command will allow you to monitor the training process by viewing the training and validation loss curves. You can view the curves by opening a web browser and going to the IP address of the training machine at port 6006. It is possible that someone else is using port 6006 if you are on a shared machine, in which case you can change the port number in the command above and use that port number to view the curves in the browser.

Your terminal will look something like this:

An alternative: train with the container shell

On a workspace that has not been migrated to airfield, open the container shell and run the training script from the workspace root, where the package lives under src/:

./container shell
python3 src/av_imitation/av_imitation/src/main.py --model cnn --dataset outdoor_v1 --epochs 1000 --batch-size 4

Start tensorboard the same way, from a second container shell.

Monitor the training process by watching the tensorboard logs and the output of the training script in the terminal. When training is complete, you'll see a message like this:

Epoch [30/1000] Train Loss: 0.0047 Val Loss: 0.0128
Saved best model to /home/orin/roboracer_ws/data/experiments/20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s/models/best_model.pth
Epoch [31/1000] Train Loss: 0.0047 Val Loss: 0.0130
Epoch [32/1000] Train Loss: 0.0047 Val Loss: 0.0129
Epoch [33/1000] Train Loss: 0.0047 Val Loss: 0.0129
Epoch [34/1000] Train Loss: 0.0047 Val Loss: 0.0129
Epoch [35/1000] Train Loss: 0.0047 Val Loss: 0.0130
Epoch [36/1000] Train Loss: 0.0047 Val Loss: 0.0130
Epoch [37/1000] Train Loss: 0.0047 Val Loss: 0.0130
Epoch [38/1000] Train Loss: 0.0047 Val Loss: 0.0130
Epoch [39/1000] Train Loss: 0.0047 Val Loss: 0.0129
Epoch [40/1000] Train Loss: 0.0046 Val Loss: 0.0128
Early stopping triggered after 40 epochs.
Training finished.

Lab Notebook

Record the number of epochs until your model converged. Also include a screenshot of the training and validation loss curves.

4. Check the Training Curves

Check the training curves to ensure the model converged.

In order to understand how well the training process worked, we will first examine the training curves. If these curves look reasonable, we will then roll out the policy on the car and evaluate its performance.

With tensorboard running (Step 3), open your browser to your development computer's IP address at port 6006, e.g. http://10.1.0.112:6006.

The following screenshot shows three different model architectures trained on the same dataset:

Ideally, the curves should show the training loss continuing to decrease while the validation loss increases.

We are using fast early stopping to prevent overfitting. So, even though you've set the number of epochs to a high value, the training will automatically stop when the validation loss stops decreasing and we'll use the model weights with the lowest validation loss.

5. Check the model's output on the validation set and convert to onnx

Open the model in the webapp's "5. Inference" tab, choose the same dataset that you trained the model on and the validation set, then click "Run Inference."

Then, to convert the model to onnx format for inference on the car, open the model in the webapp's "5. Inference" tab, choose the same dataset that you trained the model on and the validation set, then click the "Convert to ONNX" button.

After converting to ONNX, your experiment should show the text “ONNX Available”.

Alternatively, run the export through airfield, replacing 20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s with the name of your experiment:

cd ~/roboracer_ws
airfield package cmd av_imitation -- ros2 run av_imitation export_onnx --experiment-name 20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s

The ONNX export and the webapp's "Run Inference" tab use onnxruntime's CPU execution provider on the car — there is no GPU onnxruntime wheel for the Orin's JetPack yet. That is fine: the models are small, and the car runs the policy at the camera's frame rate. The comments in roboracer_ws/dependencies/arm64/onnxruntime.yaml explain why, and how to switch to a GPU build once one exists.

An alternative: export with the container shell

On a workspace that has not been migrated to airfield:

./container shell
ros2 run av_imitation export_onnx --experiment-name 20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s

6. Rollout the policy on the car

Note: If you trained the policy on a different machine, you will need to copy the onnx file to the car. For example, use this rsync command, run on the training machine, replacing XX with your car's ID: rsync -avz ~/roboracer_ws/data/experiments/20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s orin@orinXX:roboracer_ws/data/experiments/20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s. If this command errors out because mkdir fails, first ssh into the car and create the directory with the command: ssh orin@orinXX mkdir -p roboracer_ws/data/experiments/

The node that inferences the policy is ros2 run av_imitation play --experiment-name <your experiment>. It is incorporated into the imitation_inference plan, which also launches the recorder, the camera, joystick, vesc driver, robot description, GUI, VNC desktop, and rviz2. The plan reads the experiment name from the EXPERIMENT environment variable, so pass it on the launch line:

Note: Replace 20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s with the name of your experiment

cd ~/roboracer_ws
EXPERIMENT=20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s airfield project up imitation_inference

Launch from your ssh shell, not from inside tmux

EXPERIMENT is expanded by the pane's shell when the plan starts, and tmux panes inherit their environment from the running tmux server — not from the command you just typed. If you launch the plan from inside an existing tmux session, the play node may fall back to the default experiment no matter what you passed. Run airfield project down to tear down any previous session, then launch from a plain ssh shell. With no EXPERIMENT set, the plan falls back to a default experiment name.

Then, wait for the camera to load, connect the joystick and then to let the car run in autonomous mode, hold the R1 button on the joystick.

When you are done, tear the session down:

airfield project down

An alternative: roll out with tmuxinator

The car previously launched the inference nodes from a tmuxinator configuration inside the container shell. This flow still works if airfield is unavailable, and it is the only flow available on a car that has not been migrated.

cd ~/roboracer_ws/tmux/imitation_inference
export EXPERIMENT=20251215_033702_cnn_sidewalk_grey_25min_5hz_h-5s_f-2s
tmuxinator

7. Collect additional data

DAgger collects more data by letting your policy drive and recording what you do when it makes a mistake. Launch the imitation_inference plan as in step 6 and hold R1 to let the policy drive. Press the red circle on the joystick to start recording. When the car starts to drift off course, release R1 and hold L1 to drive it back on course yourself, then release L1 and hold R1 again to hand control back to the policy. Press the red circle again to stop recording.

Note: Correct the car with L1, not by moving the sticks while holding R1. The sticks will steer the car either way, but the annotator only keeps the sections where L1 was held.

When you process the bag in the annotator, the default cuts keep only the sections where L1 was held, which are your corrections. Add these bags to your dataset alongside your original demonstrations and retrain.

Troubleshooting

GPU Out of Memory

If you run into a GPU out of memory error (example error shown below), try reducing the batch size.

$ python3 av_imitation/src/main.py --model mlp --dataset outdoor_v1 --epochs 1000 --batch-size 8
Starting experiment: 20251124_235734_mlp_outdoor_v1
Using device: cuda
Loading dataset metadata...
...
Output Steps: 2
...
NvMapMemAllocInternalTagged: 1075072515 error 12
NvMapMemHandleAlloc: error 0
NvMapMemAllocInternalTagged: 1075072515 error 12
NvMapMemHandleAlloc: error 0
NvMapMemAllocInternalTagged: 1075072515 error 12
NvMapMemHandleAlloc: error 0
NvMapMemAllocInternalTagged: 1075072515 error 12
NvMapMemHandleAlloc: error 0
NvMapMemAllocInternalTagged: 1075072515 error 12
NvMapMemHandleAlloc: error 0
Traceback (most recent call last):
  File "/home/orin/workspace/src/av_imitation/av_imitation/src/main.py", line 147, in <module>
    train(args)
  File "/home/orin/workspace/src/av_imitation/av_imitation/src/main.py", line 88, in train
    optimizer.step()
  File "/usr/local/lib/python3.12/dist-packages/torch/optim/optimizer.py", line 487, in wrapper
    out = func(*args, **kwargs)
  File "/usr/local/lib/python3.12/dist-packages/torch/optim/optimizer.py", line 90, in _use_grad
    ret = func(self, *args, **kwargs)
  File "/usr/local/lib/python3.12/dist-packages/torch/optim/adam.py", line 216, in step
    has_complex = self._init_group(
  File "/usr/local/lib/python3.12/dist-packages/torch/optim/adam.py", line 160, in _init_group
    state["exp_avg_sq"] = torch.zeros_like(
RuntimeError: NVML_SUCCESS == r INTERNAL ASSERT FAILED at "/opt/pytorch/pytorch/c10/cuda/CUDACachingAllocator.cpp":838, please report a bug to PyTorch.