add g_tol parameter to gradient modes

Author: Giovanni3A

src/optimizers/gradient.jl

@@ -34,6 +34,7 @@ function train_with_gradient!(
     verbose = get(params, :verbose, true)
     compute_cost_every = get(params, :compute_cost_every, 1)
     time_limit = get(params, :time_limit, Inf)
+    g_tol = get(params, :g_tol, 0)
 
     # init parameters
     start_time = time()
@@ -87,6 +88,17 @@ function train_with_gradient!(
 
         # check time limit reach
         if time() - start_time > time_limit
+            if verbose
+                println("Time limit reached.")
+            end
+            break
+        end
+
+        # check gradient tolerance
+        if maximum(abs.(dC)) < g_tol
+            if verbose
+                println("Gradient tolerance reached.")
+            end
             break
         end
 

src/optimizers/gradient_mpi.jl

@@ -17,14 +17,15 @@ function train_with_gradient_mpi!(
     compute_cost_every = get(params, :compute_cost_every, 1)
     mpi_finalize = get(params, :mpi_finalize, true)
     time_limit = get(params, :time_limit, Inf)
+    g_tol = get(params, :g_tol, 0)
 
     JQM.mpi_init()
 
     # init parameters
     start_time = time()
     is_done = false
     best_C = Inf
-    best_θ = []
+    best_θ = extract_params(model.forecast)
     curr_C = 0.0
     trace = Array{Float64}(undef, epochs)
     dCdz = Vector{Float32}(undef, size(model.policy_vars, 1))
@@ -121,6 +122,17 @@ function train_with_gradient_mpi!(
 
             # check time limit reach
             if time() - start_time > time_limit
+                if verbose
+                    println("Time limit reached.")
+                end
+                break
+            end
+
+            # check gradient tolerance
+            if maximum(abs.(dCdy)) < g_tol
+                if verbose
+                    println("Gradient tolerance reached.")
+                end
                 break
             end