GitHub CI workflows

.github/workflows/dry-run.yml

@@ -0,0 +1,45 @@
+name: "Dry Run"
+on:
+  pull_request:
+  push:
+jobs:
+  run_steps_from_readme:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: actions/cache@v3
+      with:
+        path: |
+          nn/target
+        key: ${{ runner.os }}-cache
+    - uses: cachix/install-nix-action@v27
+      with:
+        nix_path: nixpkgs=channel:nixos-unstable
+
+    - name: Compile Rust Code
+      run: |
+        echo "Compiling Rust Code"
+        pushd nn
+        nix develop . --command cargo build --release
+        popd
+
+    - name: Generate Optimized Trajectories
+      run: |
+        echo "Generating Optimised Trajectories"
+        pushd generate_data
+        nix develop . --command env GENERATE_DATA_N_TRAJECTORIES=100 uv run quadcopter.py
+        popd
+        
+    - name: Run Training
+      run: |
+        echo "Running Training"
+        pushd nn
+        nix develop . --command env TRAIN_USE_NDARRAY=1 cargo run --release --bin train
+        popd
+        
+    - name: Run Inference
+      run: |
+        echo "Running Inference"
+        pushd nn
+        nix develop . --command env INFERENCE_USE_NDARRAY=1 cargo run --release --bin vis_results
+        popd

README.md

@@ -1,5 +1,7 @@
 # Quadcopter Optimal Landing: Neural Network Controller
 
+![GitHub CI](https://github.com/lancelet/nn-landing-poc/actions/workflows/dry-run.yml/badge.svg)
+
 WARNING: This is an initial proof-of-concept, not a carefully polished 
 project!
 

generate_data/flake.nix

@@ -6,17 +6,22 @@
     flake-utils.url = "github:numtide/flake-utils";
   };
 
-  outputs = { self, nixpkgs, flake-utils }:
-    flake-utils.lib.eachDefaultSystem (system:
-      let
-        pkgs = import nixpkgs { inherit system; };
-      in
-      {
+  outputs = {
+    self,
+    nixpkgs,
+    flake-utils,
+  }:
+    flake-utils.lib.eachDefaultSystem (
+      system: let
+        pkgs = import nixpkgs {inherit system;};
+      in {
         devShells.default = pkgs.mkShell {
           buildInputs = with pkgs; [
+            stdenv
             python312
             uv
           ];
+          LD_LIBRARY_PATH = "${pkgs.stdenv.cc.cc.lib}/lib";
         };
       }
     );

generate_data/quadcopter.py

@@ -21,10 +21,10 @@
 #---- Constants ---------------------------------------------------------------
 
 # Number of processes to use for data generation.
-N_PROCESSES = 8
+N_PROCESSES = int(os.getenv("GENERATE_DATA_N_PROCESSES", 8))
 
 # Number of trajectories to generate.
-N_TRAJECTORIES = 15000
+N_TRAJECTORIES = int(os.getenv("GENERATE_DATA_N_TRAJECTORIES", 15000))
 
 # Current project directory (based on the location of this file).
 PROJECT_DIR = Path(os.path.dirname(os.path.abspath(__file__))).resolve()

nn/Cargo.toml

@@ -5,7 +5,7 @@ edition = "2021"
 
 [dependencies]
 anyhow = "1.0.89"
-burn = { version = "0.14.0", features = ["dataset", "train", "wgpu"] }
+burn = { version = "0.14.0", features = ["dataset", "train", "wgpu", "ndarray"] }
 ndarray = "0.16.1"
 ndarray-npy = "0.9.1"
 ode_solvers = "0.4.0"

nn/flake.nix

@@ -15,23 +15,36 @@
       system: let
         pkgs = import nixpkgs {inherit system;};
         naersk-lib = pkgs.callPackage naersk {};
+        isMacOS = system == "x86_64-darwin" || system == "aarch64-darwin";
+        isLinux = system == "x86_64-linux" || system == "aarch64-linux" || system == "i686-linux";
       in {
         defaultPackage = naersk-lib.buildPackage ./.;
         devShell = with pkgs;
           mkShell {
-            buildInputs = [
-              cargo
-              darwin.apple_sdk.frameworks.IOKit
-              darwin.apple_sdk.frameworks.QuartzCore
-              ffmpeg
-              iconv
-              pre-commit
-              python3
-              python3Packages.matplotlib
-              rustc
-              rustfmt
-              rustPackages.clippy
-            ];
+            buildInputs =
+              [
+                cargo
+                ffmpeg
+                iconv
+                pre-commit
+                python3
+                python3Packages.matplotlib
+                rustc
+                rustfmt
+                rustPackages.clippy
+              ]
+              ++ lib.optionals isMacOS [
+                darwin.apple_sdk.frameworks.IOKit
+                darwin.apple_sdk.frameworks.QuartzCore
+              ]
+              ++ lib.optionals isLinux [
+                vulkan-loader
+                vulkan-headers
+                vulkan-validation-layers
+                vulkan-tools
+                glslang
+                shaderc
+              ];
             RUST_SRC_PATH = rustPlatform.rustLibSrc;
           };
       }

nn/src/bin/main_train.rs

@@ -1,12 +1,20 @@
+use std::env;
+
 use train_quadcopter::training;
 
 use anyhow::Result;
+use burn::backend::ndarray::{NdArray, NdArrayDevice};
 use burn::backend::wgpu::{Wgpu, WgpuDevice};
 use burn::backend::Autodiff;
 
 fn main() -> Result<()> {
     println!("Quadcopter Training");
 
-    let device = WgpuDevice::default();
-    training::run::<Autodiff<Wgpu>>(device)
+    if env::var("TRAIN_USE_NDARRAY").is_ok() {
+        let device = NdArrayDevice::default();
+        training::run::<Autodiff<NdArray>>(device)
+    } else {
+        let device = WgpuDevice::default();
+        training::run::<Autodiff<Wgpu>>(device)
+    }
 }

nn/src/bin/main_vis_results.rs

@@ -1,10 +1,13 @@
+use std::env;
 use std::f32::consts::PI;
 use std::fs::create_dir_all;
 
 use anyhow::{bail, Result};
+use burn::backend::ndarray::NdArrayDevice;
 use burn::backend::wgpu::WgpuDevice;
-use burn::backend::Wgpu;
+use burn::backend::{NdArray, Wgpu};
 use burn::module::Module;
+use burn::prelude::Backend;
 use burn::record::{DefaultFileRecorder, FullPrecisionSettings};
 
 use plotpy::{Curve, Plot};
@@ -21,10 +24,18 @@ fn main() -> Result<()> {
     println!("Trained Model Visualization");
 
     // Load the model weights from the trained file.
-    let device = WgpuDevice::default();
+    if env::var("INFERENCE_USE_NDARRAY").is_ok() {
+        let device = NdArrayDevice::default();
+        run_inference::<NdArray>(device)
+    } else {
+        let device = WgpuDevice::default();
+        run_inference::<Wgpu>(device)
+    }
+}
+
+fn run_inference<B: Backend>(device: B::Device) -> Result<()> {
     let file_path = "./burn-training-artifacts/model.mpk";
-    let config = ModelConfig::new();
-    let model = config.init::<Wgpu>(&device);
+    let model = ModelConfig::new().init::<B>(&device);
     let model = model.load_file(
         file_path,
         &DefaultFileRecorder::<FullPrecisionSettings>::new(),

nn/src/training.rs

@@ -1,3 +1,5 @@
+use std::io::IsTerminal;
+
 use crate::{
     data::{load_datasets, ExampleBatcher, MeanStdev, StateMeanStdev},
     model::ModelConfig,
@@ -14,6 +16,7 @@ use burn::{
             store::{Aggregate, Direction, Split},
             LossMetric,
         },
+        renderer::MetricsRenderer,
         LearnerBuilder, MetricEarlyStoppingStrategy, StoppingCondition,
     },
 };
@@ -65,7 +68,7 @@ pub fn run<B: AutodiffBackend>(device: B::Device) -> Result<()> {
         .build(test_dataset);
 
     // Model
-    let learner = LearnerBuilder::new(ARTIFACT_DIR)
+    let learner_builder = LearnerBuilder::new(ARTIFACT_DIR)
         .metric_train_numeric(LossMetric::new())
         .metric_valid_numeric(LossMetric::new())
         .with_file_checkpointer(DefaultFileRecorder::<FullPrecisionSettings>::new())
@@ -76,9 +79,13 @@ pub fn run<B: AutodiffBackend>(device: B::Device) -> Result<()> {
             StoppingCondition::NoImprovementSince { n_epochs: 1 },
         ))
         .devices(vec![device.clone()])
-        .num_epochs(num_epochs)
-        .summary()
-        .build(model, optimizer.init(), lr);
+        .num_epochs(num_epochs);
+    let learner_builder = if !std::io::stdout().is_terminal() {
+        learner_builder.renderer(StdoutMetricsRenderer)
+    } else {
+        learner_builder
+    };
+    let learner = learner_builder.summary().build(model, optimizer.init(), lr);
 
     let model_trained = learner.fit(dataloader_train, dataloader_test);
 
@@ -91,3 +98,19 @@ pub fn run<B: AutodiffBackend>(device: B::Device) -> Result<()> {
 
     Ok(())
 }
+
+struct StdoutMetricsRenderer;
+
+impl MetricsRenderer for StdoutMetricsRenderer {
+    fn update_train(&mut self, _state: burn::train::renderer::MetricState) {}
+
+    fn update_valid(&mut self, _state: burn::train::renderer::MetricState) {}
+
+    fn render_train(&mut self, item: burn::train::renderer::TrainingProgress) {
+        println!("Training iteration: {}", item.iteration);
+    }
+
+    fn render_valid(&mut self, item: burn::train::renderer::TrainingProgress) {
+        println!("Validation iteration: {}", item.iteration);
+    }
+}

walkthrough.md

@@ -117,3 +117,15 @@ to a movie with `ffmpeg` like this:
 ```bash
 ffmpeg -framerate 60 -i ./animation/%05d.png -vcodec libx264 -s 1080x1080 -pix_fmt yuv420p animation.mov
 ```
+
+## CI Build
+
+There is a CI build which runs the project through its basic steps. See
+`.github/workflows/dry-run.yml`. The main differences between the CI build
+and running the project locally are:
+  - The CI build fetches some dependencies using Nix.
+  - Only 100 training trajectories are generated.
+  - The renderer for the fancy training GUI is replaced by one that dumps
+    progress to `stdout`.
+  - The `NDArray` backend is used instead of `wgpu`. I couldn't figure out
+    the trinity of Nix, WGPU and GitHub.