Merge pull request #7 from rejuvyesh/jkg/folds

Allow using Folds to parallelize `ARS`

Author: thowell

Project.toml

@@ -35,11 +35,11 @@ JLD2 = "0.4.21"
 LaTeXStrings = "1.3.0"
 LightGraphs = "1.3.5"
 LightXML = "<0.9.0, 0.9.0"
-MeshCat = "0.13.0 - 0.13.0"
+MeshCat = "0.14"
 Meshing = "0.5.7"
 Parameters = "0.12"
 Polyhedra = "0.6.18"
 Rotations = "1.0.2 - 1.0.2"
 Scratch = "1.1"
-StaticArrays = "0.12, 1.0"
+StaticArrays = "1.4"
 julia = "1.6"

examples/Project.toml

@@ -1,12 +1,14 @@
 [deps]
 Dojo = "ac60b53e-8d92-4c83-b960-e78698fa1916"
+Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 IterativeLQR = "605048dd-e178-462b-beb9-98a09398ef27"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 PGFPlots = "3b7a836e-365b-5785-a47d-02c71176b4aa"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [compat]
+Folds = "0.2"
 IterativeLQR = "0.1.1"
 JLD2 = "0.4.21"
 Literate = "2.13.0"

examples/reinforcement_learning/algorithms/ars.jl

@@ -3,12 +3,13 @@
 ################################################################################
 using LinearAlgebra
 using Statistics
+using Folds
 
 import LinearAlgebra.normalize
-import GeometryBasics.update
+import Dojo.GeometryBasics.update
 
 # ARS options: hyper parameters
-@with_kw struct HyperParameters{T}
+Base.@kwdef struct HyperParameters{T}
     main_loop_size::Int = 100
     horizon::Int = 200
     step_size::T = 0.02
@@ -114,16 +115,19 @@ function rollout_policy(θ::Matrix, env::Environment, normalizer::Normalizer, hp
 end
 
 function train(env::Environment, policy::Policy{T}, normalizer::Normalizer{T},
-        hp::HyperParameters{T}; distributed=false) where T
+        hp::HyperParameters{T}; distributed=false, usefolds=false, foldsexec=Folds.ThreadedEx(;basesize=1)) where T
     println("Training linear policy with Augmented Random Search (ARS)\n ")
     if distributed
         envs = [deepcopy(env) for i = 1:(2 * hp.n_directions)]
         normalizers = [deepcopy(normalizer) for i = 1:(2 * hp.n_directions)]
         hps = [deepcopy(hp) for i = 1:(2 * hp.n_directions)]
         print("  $(nprocs()) processors")
+    elseif usefolds
+        envs = [deepcopy(env) for i = 1:(2*hp.n_directions)]
+        print("  $(Threads.nthreads()) threads with Folds")
     else
-        envs = [deepcopy(env) for i = 1:Threads.nthreads()]
-        print("  $(Threads.nthreads()) threads")
+        envs = [deepcopy(env) for i = 1:(Threads.nthreads())]
+        print(" $(Threads.nthreads()) ")
     end
 
     # pre-allocate for rewards
@@ -134,14 +138,20 @@ function train(env::Environment, policy::Policy{T}, normalizer::Normalizer{T},
         θs, δs = sample_policy(policy)
 
         # evaluate policies
-        if distributed
-            rewards .= pmap(rollout_policy, θs, envs, normalizers, hps)
-        else
-            Threads.@threads for k = 1:(2 * hp.n_directions)
-                rewards[k] = rollout_policy(θs[k], envs[Threads.threadid()], normalizer, hp)
+        roll_time = @elapsed begin
+            if distributed
+                rewards .= pmap(rollout_policy, θs, envs, normalizers, hps)
+            elseif usefolds
+                @assert length(envs) == size(θs, 1) "$(length(envs))"
+                Folds.map!(rewards, θs, envs, foldsexec) do θ, env
+                    rollout_policy(θ, env, normalizer, hp)
+                end
+            else
+                Threads.@threads for k = 1:(2 * hp.n_directions)
+                    rewards[k] = rollout_policy(θs[k], envs[Threads.threadid()], normalizer, hp)
+                end
             end
         end
-
         # reward evaluation
         r_max = [max(rewards[k], rewards[hp.n_directions + k]) for k = 1:hp.n_directions]
         σ_r = std(rewards)
@@ -152,7 +162,7 @@ function train(env::Environment, policy::Policy{T}, normalizer::Normalizer{T},
         update(policy, rollouts, σ_r)
 
         # finish, print:
-        println("episode $episode reward_evaluation $(mean(rewards))")
+        println("episode $episode reward_evaluation $(mean(rewards)). Took $(roll_time) seconds")
     end
 
     return nothing