train_naive_point.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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using CUDA\n",
    "\n",
    "# Check if GPU is available\n",
    "if CUDA.functional()\n",
    "    device!(1)\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": [
    "include(\"dataset.jl\")\n",
    "X = Hyperrectangle(low = [0, 0, -1, -1], high = [4,4, 1, 1])\n",
    "U = Hyperrectangle(low = [-1, -1], high = [1,1])\n",
    "X_unsafe = Hyperrectangle(low = [1.5, 0, -1, -1], high = [2.5,2, 1, 1])\n",
    "\n",
    "A = [0. 0 1 0;\n",
    "    0 0 0 1;\n",
    "    0 0 0 0;\n",
    "    0 0 0 0;]\n",
    "B = [0. 0;\n",
    "    0 0;\n",
    "    1 0;\n",
    "    0 1;]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_training_data = load_object(\"point_seq_training_data.jld2\")\n",
    "raw_test_data = load_object(\"point_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(4 => 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",
    "μ = 0.1\n",
    "α = 0.0\n",
    "use_pgd=true # always set true to train a policy-independent CBF\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.NADAM(ini_lr, (0.9, 0.999), 0.1), model)  # will store optimiser momentum, etc.\n",
    "sched = ParameterSchedulers.Stateful(Step(ini_lr, lr_decay_rate, lr_decay_epoch)) # setup schedule of your choice\n",
    "\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",
    "        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, α=α) + μ .* 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, α=α) + μ .* loss_regularization(model, x_batch, y_init_batch)\n",
    "        push!(training_loss_epcoh, loss)  # logging, outside gradient context\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",
    "        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, α=α) + μ .* 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(\"naive_model_1_0_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": [
    "\n",
    "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
}