train_adv_planar_quad.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Revise\n",
    "using LazySets\n",
    "using DifferentialEquations\n",
    "using LazySets\n",
    "using ProgressBars\n",
    "using JLD2\n",
    "using Flux\n",
    "using LinearAlgebra\n",
    "using Zygote\n",
    "using ReverseDiff\n",
    "using Plots\n",
    "using Statistics\n",
    "using Optimisers, ParameterSchedulers\n",
    "using RobotDynamics\n",
    "using RobotZoo\n",
    "using Random\n",
    "import RobotDynamics as RD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using CUDA\n",
    "\n",
    "# Check if GPU is available\n",
    "if CUDA.functional()\n",
    "    device!(2)\n",
    "    CUDA.allowscalar(false)  # Disallow scalar operations on the GPU (optional)\n",
    "else\n",
    "    println(\"GPU is not available. Using CPU.\")\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = Hyperrectangle(low = [0, 0, -0.1, -1, -1 ,-1], high = [4,4, 0.1, 1,1,1])\n",
    "U = Hyperrectangle(low = [4, 4], high = [6,6])\n",
    "X_unsafe = Hyperrectangle(low = [1.5, 0,-0.1,-1, -1 ,-1], high = [2.5,2, 0.1, 1,1,1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_training_data = load_object(\"planar_quadrotor_seq_training_data.jld2\")\n",
    "raw_test_data = load_object(\"planar_quadrotor_seq_test_data.jld2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "include(\"affine_dynamics.jl\")\n",
    "using Flux\n",
    "using ProgressBars\n",
    "using JLD2\n",
    "\n",
    "batchsize = 128\n",
    "@show sum(raw_training_data[3,:]), size(raw_training_data[3,:])\n",
    "state_dim = size(reduce(hcat,raw_training_data[1,:]))[1]\n",
    "\n",
    "training_data = raw_training_data\n",
    "test_data = raw_test_data\n",
    "\n",
    "model = Chain(\n",
    "    Dense(6 => 16, relu),   # activation function inside layer\n",
    "    Dense(16 => 64, relu),   # activation function inside layer\n",
    "    Dense(64 => 16, relu),   # activation function inside layer\n",
    "    Dense(16 => 1)\n",
    ")\n",
    "\n",
    "# model = Chain(\n",
    "#     Dense(4 => 16, sigmoid),   # activation function inside layer\n",
    "#     Dense(16 => 64, sigmoid),   # activation function inside layer\n",
    "#     Dense(64 => 16, sigmoid),   # activation function inside layer\n",
    "#     Dense(16 => 1)\n",
    "# )\n",
    "\n",
    "# model = Chain(\n",
    "#     Dense(4 => 16, tanh_fast),   # activation function inside layer\n",
    "#     Dense(16 => 64, tanh_fast),   # activation function inside layer\n",
    "#     Dense(64 => 16, tanh_fast),   # activation function inside layer\n",
    "#     Dense(16 => 1)\n",
    "# )\n",
    "\n",
    "train_loader = Flux.DataLoader(training_data, batchsize=batchsize, shuffle=true)\n",
    "test_loader = Flux.DataLoader(test_data, batchsize=batchsize, shuffle=true)\n",
    "\n",
    "λ = 1\n",
    "μ = 1\n",
    "α = 0.0\n",
    "use_pgd=true\n",
    "use_adv = true\n",
    "lr_lambda = 0.0001\n",
    "mu = 0.9\n",
    "V_lambda=0\n",
    "\n",
    "lr_decay_rate = 0.2\n",
    "lr_decay_epoch =4\n",
    "total_epoch = 20\n",
    "\n",
    "ini_lr = 0.01\n",
    "optim = Flux.setup(Flux.Optimise.AdamW(ini_lr, (0.9, 0.999), 0.001), model)\n",
    "sched = ParameterSchedulers.Stateful(Step(ini_lr, lr_decay_rate, lr_decay_epoch)) # setup schedule of your choice\n",
    "\n",
    "eps = 1e-3\n",
    "dyn_model = RobotZoo.PlanarQuadrotor()\n",
    "n,m = RD.dims(dyn_model)\n",
    "\n",
    "training_losses = []\n",
    "test_losses = []\n",
    "least_loss = 1000\n",
    "test_loss = 0\n",
    "loss = 0\n",
    "for epoch in ProgressBar(1:total_epoch)\n",
    "    training_loss_epcoh = []\n",
    "    test_loss_epcoh = []\n",
    "    ∇l_lambda = 0\n",
    "    ∇l_alpha = 0\n",
    "    for item in train_loader\n",
    "        ∇l_alpha = 0\n",
    "        ∇l_lambda = 0\n",
    "        x_batch = reduce(hcat,item[1,:])\n",
    "        u_batch = reduce(hcat,item[2,:])\n",
    "        y_init_batch = reduce(hcat,item[3,:])\n",
    "        A = []\n",
    "        B = []\n",
    "        Δ = []\n",
    "        for i in 1:size(x_batch, 2)\n",
    "            z = RD.KnotPoint(x_batch[:, i],u_batch[:, i],0.0,1e-3 ) \n",
    "            ∇f = zeros(n, n + m)\n",
    "            RD.jacobian!(RD.StaticReturn(), RD.ForwardAD(), dyn_model, ∇f, zeros(n), z)\n",
    "            A_ = ∇f[:, 1:n]\n",
    "            B_ = ∇f[:, n+1:end]\n",
    "            Δ_ = RobotDynamics.dynamics(dyn_model, x_batch[:, i] .- eps, u_batch[:, i].-eps) - A_ * (x_batch[:, i].-eps) - B_ * (u_batch[:, i] .- eps)\n",
    "            push!(A, A_)\n",
    "            push!(B, B_)\n",
    "            push!(Δ, Δ_)\n",
    "        end\n",
    "        A = cat(A..., dims=3)\n",
    "        B = cat(B..., dims=3)\n",
    "        Δ = cat(Δ..., dims=2)\n",
    "        use_pgd && (u_batch = pgd_find_u_notce(model, A, x_batch, B, u_batch, U; α = α,Δ=Δ))\n",
    "        if use_adv\n",
    "            X_lcoal = [Hyperrectangle(x_batch[:, i], radius_hyperrectangle(X) ./ 20) for i=1:size(x_batch)[2]]\n",
    "            x_batch = pgd_find_x_notce(model, A, x_batch, B, u_batch, X_lcoal; α = α,Δ=Δ)\n",
    "        else\n",
    "            x_batch = x_batch\n",
    "        end\n",
    "        use_pgd && (u_batch = pgd_find_u_notce(model, A, x_batch, B, u_batch, U; α = α,Δ=Δ))\n",
    "        training_loss, grads = Flux.withgradient(model) do m \n",
    "            loss_naive_safeset(m, x_batch, y_init_batch) + λ .* loss_naive_fi(m, A, x_batch, B, u_batch,y_init_batch;use_pgd=false,use_adv = false, α=α,Δ=Δ) + μ .* loss_regularization(m, x_batch, y_init_batch)\n",
    "        end\n",
    "        # Update the parameters so as to reduce the objective,\n",
    "        # according the chosen optimisation rule:\n",
    "        Flux.update!(optim, model, grads[1])\n",
    "\n",
    "        loss = loss_naive_safeset(model, x_batch, y_init_batch) + λ .* loss_naive_fi(model, A, x_batch, B, u_batch,y_init_batch;use_pgd=use_pgd,use_adv = false, α=α,Δ=Δ) + μ .* loss_regularization(model, x_batch, y_init_batch)\n",
    "        push!(training_loss_epcoh, loss)  # logging, outside gradient context\n",
    "        \n",
    "    end\n",
    "    for item in test_loader\n",
    "        x_batch = reduce(hcat,item[1,:])\n",
    "        u_batch = reduce(hcat,item[2,:])\n",
    "        y_init_batch = reduce(hcat,item[3,:])\n",
    "        \n",
    "        A = []\n",
    "        B = []\n",
    "        Δ = []\n",
    "        for i in 1:size(x_batch, 2)\n",
    "            z = RD.KnotPoint(x_batch[:, i],u_batch[:, i],0.0,1e-3 ) \n",
    "            ∇f = zeros(n, n + m)\n",
    "            RD.jacobian!(RD.StaticReturn(), RD.ForwardAD(), dyn_model, ∇f, zeros(n), z)\n",
    "            A_ = ∇f[:, 1:n]\n",
    "            B_ = ∇f[:, n+1:end]\n",
    "            Δ_ = RobotDynamics.dynamics(dyn_model, x_batch[:, i] .- eps, u_batch[:, i].-eps) - A_ * (x_batch[:, i].-eps) - B_ * (u_batch[:, i] .- eps)\n",
    "            push!(A, A_)\n",
    "            push!(B, B_)\n",
    "            push!(Δ, Δ_)\n",
    "        end\n",
    "        A = cat(A..., dims=3)\n",
    "        B = cat(B..., dims=3)\n",
    "        Δ = cat(Δ..., dims=2)\n",
    "\n",
    "        test_loss =  loss_naive_safeset(model, x_batch, y_init_batch) + λ .* loss_naive_fi(model, A, x_batch, B, u_batch,y_init_batch;use_pgd=use_pgd,use_adv = false, α=α,Δ=Δ) + μ .* loss_regularization(model, x_batch, y_init_batch)\n",
    "        push!(test_loss_epcoh, test_loss)\n",
    "    end\n",
    "    nextlr = ParameterSchedulers.next!(sched) # advance schedule\n",
    "    Optimisers.adjust!(optim, nextlr) # update optimizer state, by default this changes the learning rate `eta`\n",
    "    lr_lambda = lr_lambda * lr_decay_rate^(floor(epoch / lr_decay_epoch))\n",
    "    @show epoch, loss, test_loss\n",
    "    model_state = Flux.state(model)\n",
    "    jldsave(\"planar_quad_wd0.001_adv20_model_1_0_1_pgd_relu_$epoch.jld2\"; model_state)\n",
    "    push!(training_losses, sum(training_loss_epcoh)) \n",
    "    push!(test_losses, sum(test_loss_epcoh))\n",
    "\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots\n",
    "\n",
    "function plot_loss(train_loss::Vector, test_loss::Vector; xlabel=\"Epoch\", ylabel=\"Loss\", title=\"Training and Test Loss\")\n",
    "    plot(train_loss ./ (size(training_data[3,:] / batchsize)), label=\"Training\", xlabel=xlabel, ylabel=ylabel, title=title)\n",
    "    plot!(test_loss ./ (size(test_data[3,:] / batchsize)), label=\"Test\")\n",
    "end\n",
    "\n",
    "plot_loss(training_losses, test_losses)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.4",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}