From 029f7f29f2bd7765bc205dbdafeec308c01925af Mon Sep 17 00:00:00 2001 From: MoBit Date: Sun, 12 Dec 2021 15:29:08 +0100 Subject: [PATCH] Restructure for RL --- src/analysis/pair_correlation_function.jl | 2 +- src/animation.jl | 2 +- src/benchmark.jl | 6 +- src/data.jl | 32 +++- src/reinforcement_learning.jl | 182 ++++++++++++++++++---- src/run.jl | 117 +++++++------- src/setup.jl | 98 ++++++++---- src/simulation.jl | 32 +++- src/utils.jl | 16 -- 9 files changed, 339 insertions(+), 148 deletions(-) delete mode 100644 src/utils.jl diff --git a/src/analysis/pair_correlation_function.jl b/src/analysis/pair_correlation_function.jl index b27a921..bf17fb9 100644 --- a/src/analysis/pair_correlation_function.jl +++ b/src/analysis/pair_correlation_function.jl @@ -12,7 +12,7 @@ function plot_g(radius, g, variables) yticks=0:ceil(Int64, maximum(g)), xlabel=L"r", ylabel=L"g(r)", - title="v = $(variables.v)", + title="v₀ = $(variables.v₀)", ) scatterlines!(ax, radius, g; color=:white, markercolor=:red) diff --git a/src/animation.jl b/src/animation.jl index 6782332..a6428f9 100644 --- a/src/animation.jl +++ b/src/animation.jl @@ -133,7 +133,7 @@ function animate_after_loading(dir, animation_path, sim_consts, framerate, debug skin_colors = Observable(Vector{RGBAf}(undef, n_particles)) end - particle_radius = sim_consts.particle_diameter / 2 + particle_radius = sim_consts.particle_radius interaction_r = sim_consts.interaction_r bundle_paths = readdir("$dir/bundles"; join=true, sort=false) diff --git a/src/benchmark.jl b/src/benchmark.jl index 8247e13..434af75 100644 --- a/src/benchmark.jl +++ b/src/benchmark.jl @@ -8,13 +8,13 @@ using ReCo function run_benchmarks( dir::String=""; n_particles::Int64=1000, - v::Float64=20.0, + v₀::Float64=20.0, duration::Float64=2.0, n_bundle_snapshots::Int64=0, comment="", ) if length(dir) == 0 - dir = init_sim(; n_particles=n_particles, v=v, parent_dir="benchmark") + dir = init_sim(; n_particles=n_particles, v₀=v₀, parent_dir="benchmark") end benchmark = @benchmark run_sim( @@ -28,7 +28,7 @@ function run_benchmarks( f, Dict( "n_particles" => n_particles, - "v" => v, + "v₀" => v₀, "duration" => duration, "n_bundle_snapshots" => n_bundle_snapshots, "comment" => comment, diff --git a/src/data.jl b/src/data.jl index dd6e696..1698e0b 100644 --- a/src/data.jl +++ b/src/data.jl @@ -1,5 +1,7 @@ using StaticArrays: SVector using JLD2: JLD2 +using JSON3: JSON3 +using OrderedCollections: OrderedDict struct Bundle t::Vector{Float64} @@ -33,14 +35,6 @@ function save_snapshot!( return nothing end -function set_sim_state(dir::String, n_bundles::Int64, T::Float64) - open("$dir/sim_state.json", "w") do f - JSON3.write(f, (n_bundles=n_bundles, T=round(T; digits=3))) - end - - return nothing -end - function save_bundle(dir::String, bundle::Bundle, n::Int64, T::Float64) bundles_dir = "$dir/bundles" mkpath(bundles_dir) @@ -49,5 +43,27 @@ function save_bundle(dir::String, bundle::Bundle, n::Int64, T::Float64) set_sim_state(dir, n, T) + return nothing +end + +function set_sim_state(dir::String, n_bundles::Int64, T::Float64) + open("$dir/sim_state.json", "w") do f + JSON3.write(f, (n_bundles=n_bundles, T=round(T; digits=3))) + end + + return nothing +end + +function struct_to_ordered_dict(s) + return OrderedDict(key => getfield(s, key) for key in propertynames(s)) +end + +function write_struct_to_json(s, path_without_extension::String) + ordered_dict = struct_to_ordered_dict(s) + + open("$path_without_extension.json", "w") do f + JSON3.write(f, ordered_dict) + end + return nothing end \ No newline at end of file diff --git a/src/reinforcement_learning.jl b/src/reinforcement_learning.jl index 558fcf0..2f7d2d5 100644 --- a/src/reinforcement_learning.jl +++ b/src/reinforcement_learning.jl @@ -1,24 +1,45 @@ -using ReinforcementLearning +module ReCoRL -mutable struct ReCoEnvParams - n_particles::Int64 - half_box_len::Float64 - skin_r::Float64 +using ReinforcementLearning +using Intervals + +import Base: run + +struct DistanceState{L<:Bound} + interval::Interval{Float64,L,Closed} + + function DistanceState{L}(lower::Float64, upper::Float64) where {L<:Bound} + return new(Interval{Float64,L,Closed}(lower, upper)) + end +end + +struct DirectionState + interval::Interval{Float64,Open,Closed} + + function DirectionState(lower::Float64, upper::Float64) + return new(Interval{Float64,Closed,Open}(lower, upper)) + end +end + +mutable struct EnvParams action_space::Vector{Tuple{Float64,Float64}} - state_space::Vector{Tuple{Symbol,Symbol}} + distance_state_space::Vector{DistanceState} + direction_state_space::Vector{DirectionState} + state_space::Vector{Tuple{Union{DistanceState,Nothing},Union{DirectionState,Nothing}}} reward::Float64 - function ReCoEnvParams( - n_particles::Int64, - half_box_len::Float64, - skin_r::Float64, - n_v_actions::Int64, - n_ω_actions::Int64; - max_v::Float64=80.0, - max_ω::Float64=float(π), + function EnvParams( + min_distance::Float64, + max_distance::Float64; + n_v_actions::Int64=2, + n_ω_actions::Int64=3, + max_v::Float64=20.0, + max_ω::Float64=π / 1.5, + n_distance_states::Int64=3, + n_direction_states::Int64=4, ) - @assert half_box_len > 0 - @assert skin_r > 0 + @assert min_distance > 0.0 + @assert max_distance > min_distance @assert n_v_actions > 1 @assert n_ω_actoins > 1 @assert max_v > 0 @@ -39,40 +60,133 @@ mutable struct ReCoEnvParams end end - distance_state_space = (:big, :medium, :small) - direction_state_space = (:before, :behind, :left, :right) + distance_range = + min_distance:((max_distance - min_distance) / n_distance_states):max_distance - n_states = undef, length(distance_state_space) * length(direction_state_space) + 1 + distance_state_space = Vector{DistanceState}(undef, n_distance_states) - state_space = Vector{Tuple{Symbol,Symbol}}(n_states) + for i in 1:n_distance_states + if i == 1 + bound = Closed + else + bound = Open + end + + distance_state_space[i] = DistanceState{bound}( + distance_range[i], distance_range[i + 1] + ) + end + + direction_range = 0.0:(2 * π / n_direction_states):(2 * π) + + direction_state_space = Vector{DirectionState}(undef, n_direction_states) + + for i in 1:n_direction_states + direction_state_space[i] = DirectionState( + direction_range[i], direction_range[i + 1] + ) + end + + state_space = Vector{Tuple{DistanceState,DirectionState}}(undef, n_states) ind = 1 - for distance in distance_state_space - for direction in direction_state_space - state_space[ind] = (distance, direction) + for distance_state in distance_state_space + for direction_state in direction_state_space + state_space[ind] = (distance_state, direction_state) ind += 1 end end - state_space[ind] = (:none, :none) + state_space[ind] = (nothing, nothing) - return new(n_particles, half_box_len, skin_r, action_space, state_space, 0.0) + return new( + action_space, distance_state_space, direction_state_space, state_space, 0.0 + ) end end -mutable struct ReCoEnv <: AbstractEnv - params::ReCoEnvParams +mutable struct Env <: AbstractEnv + params::EnvParams particle::Particle - state::Tuple{Symbol,Symbol} + state::Tuple{Union{DistanceState,Nothing},Union{DirectionState,Nothing}} - function ReCoEnv(params::ReCoEnvParams, particle::Particle) - return new(params, particle, (:none, :none)) + function Env(params::EnvParams, particle::Particle) + return new(params, particle, (nothing, nothing)) end end -RLBase.state_space(env::ReCoEnv) = env.state_space +struct Params + envs::Vector{Env} + # agents + actions::Vector{Tuple{Float64,Float64}} + env_params::EnvParams + n_steps_before_actions_update::Int64 -RLBase.state(env::ReCoEnv) = env.state + function Params( + n_particles::Int64, env_params::EnvParams, n_steps_before_actions_update::Int64 + ) + return new( + Vector{Env}(undef, n_particles), + Vector{Tuple{Float64,Float64}}(undef, n_particles), + env_params, + n_steps_before_actions_update, + ) + end +end -RLBase.action_space(env::ReCoEnv) = env.params.action_space +RLBase.state_space(env::Env) = env.state_space -RLBase.reward(env::ReCoEnv) = env.params.reward \ No newline at end of file +RLBase.state(env::Env) = env.state + +RLBase.action_space(env::Env) = env.params.action_space + +RLBase.reward(env::Env) = env.params.reward + +function pre_integration_hook!() end + +function integration_hook() end + +function post_integration_hook() end + +function run( + n_episodes::Int64=100, + episode_duration::Float64=5.0, + update_actions_at::Float64=0.1, + n_particles::Int64=100, +) + @assert n_episodes > 0 + @assert episode_duration > 0 + @assert update_actions_at in 0.01:0.01:episode_duration + @assert episode_duration % update_actions_at == 0 + @assert n_particles > 0 + + Random.seed!(42) + + # envs + # agents + # pre_experiment + + sim_consts = gen_sim_consts(n_particles, v₀) + + env_params = EnvParams(sim_consts.particle_radius, sim_consts.skin_r) + + rl_params = Params(n_particles, env_params, n_steps_before_actions_update) + + for episode in 1:n_episodes + # reset + # pre_episode + + dir = init_sim_with_sim_consts(; sim_consts, parent_dir="RL") + + run_sim( + dir; + duration=episode_duration, + seed=rand(1:typemax(Int64)), + rl_params=rl_params, + skin_r=skin_r, + ) + end + + return nothing +end + +end # module \ No newline at end of file diff --git a/src/run.jl b/src/run.jl index 250b128..26df8af 100644 --- a/src/run.jl +++ b/src/run.jl @@ -3,18 +3,20 @@ using JSON3: JSON3 using JLD2: JLD2 using StaticArrays: SVector +empty_hook(args...) = nothing + function run_sim( dir::String; duration::Float64, snapshot_at::Float64=0.1, - n_steps_before_verlet_list_update::Int64=100, seed::Int64=42, n_bundle_snapshots::Int64=100, + rl_params::Union{ReCo.Params,Nothing}=nothing, ) @assert length(dir) > 0 @assert duration > 0 @assert snapshot_at in 0.001:0.001:duration - @assert n_steps_before_verlet_list_update >= 0 + @assert duration % snapshot_at == 0 @assert seed > 0 @assert n_bundle_snapshots >= 0 @@ -22,15 +24,9 @@ function run_sim( sim_consts = JSON3.read(read("$dir/sim_consts.json", String)) - skin_r = - 1.5 * ( - 2 * sim_consts.v * n_steps_before_verlet_list_update * sim_consts.δt + - 2 * sim_consts.interaction_r - ) - integration_steps = floor(Int64, duration / sim_consts.δt) + 1 - n_steps_before_snapshot = round(Int64, snapshot_at / sim_consts.δt) + n_steps_before_snapshot = round(Int64, snapshot_at / sim_consts.δt) n_snapshots = floor(Int64, integration_steps / n_steps_before_snapshot) + 1 n_bundle_snapshots = min(n_snapshots, n_bundle_snapshots) @@ -38,27 +34,63 @@ function run_sim( sim_state = JSON3.read(read("$dir/sim_state.json", String)) n_bundles = sim_state.n_bundles + T0::Float64 = sim_state.T + T = T0 + duration + + if n_bundle_snapshots > 0 + save_data = true + else + save_data = false + end + + @async begin + start_datetime = now() + + run_params = (; + # Input + duration, + snapshot_at, + seed, + n_bundle_snapshots, + # Calculated + integration_steps, + n_steps_before_snapshot, + n_snapshots, + T, + # Read + T0, + start_datetime, + ) + + next_bundle = n_bundles + 1 + + runs_dir = "$dir/runs" + + if save_data + write_struct_to_json(run_params, "$runs_dir/run_params_$next_bundle") + end + end + bundles_dir = "$dir/bundles" bundle = JLD2.load_object("$bundles_dir/bundle_$n_bundles.jld2") - particles = generate_particles(bundle, sim_consts.n_particles) + particles = gen_particles(bundle, sim_consts.n_particles) args = ( - v=sim_consts.v, - skin_r=skin_r, - skin_r²=skin_r^2, + v₀=sim_consts.v₀, + skin_r=sim_consts.skin_r, + skin_r²=sim_consts.skin_r^2, n_snapshots=n_snapshots, c₁=4 * sim_consts.ϵ * 6 * sim_consts.σ^6 * sim_consts.δt * sim_consts.μ, c₂=2 * sim_consts.σ^6, c₃=sqrt(2 * sim_consts.D₀ * sim_consts.δt), c₄=sqrt(2 * sim_consts.Dᵣ * sim_consts.δt), - vδt=sim_consts.v * sim_consts.δt, + v₀δt=sim_consts.v₀ * sim_consts.δt, μ=sim_consts.μ, interaction_r=sim_consts.interaction_r, interaction_r²=sim_consts.interaction_r^2, n_particles=sim_consts.n_particles, half_box_len=sim_consts.half_box_len, - particle_diameter=sim_consts.particle_diameter, particles=particles, particles_c=[particles[i].c for i in 1:(sim_consts.n_particles)], verlet_lists=[ @@ -67,45 +99,20 @@ function run_sim( ], n_bundle_snapshots=n_bundle_snapshots, bundle=Bundle(sim_consts.n_particles, n_bundle_snapshots), - box=Box(SVector(2 * sim_consts.half_box_len, 2 * sim_consts.half_box_len), skin_r), + box=Box( + SVector(2 * sim_consts.half_box_len, 2 * sim_consts.half_box_len), + sim_consts.skin_r, + ), ) - T0::Float64 = sim_state.T - T = T0 + duration - - start_datetime = now() - - run_params = (; - # Input - duration, - snapshot_at, - n_steps_before_verlet_list_update, - seed, - n_bundle_snapshots, - # Calculated - skin_r, - integration_steps, - n_steps_before_snapshot, - n_snapshots, - T, - # Read - T0, - start_datetime, - ) - - next_bundle = n_bundles + 1 - - runs_dir = "$dir/runs" - mkpath(runs_dir) - - if n_bundle_snapshots > 0 - save_data = true - - open("$runs_dir/run_params_$next_bundle.json", "w") do f - JSON3.write(f, run_params) - end + if !isnothing(rl_params) + pre_integration_hook! = ReCoRL.pre_integration_hook + integration_hook = ReCoRL.integration_hook + post_integration_hook = ReCoRL.post_integration_hook else - save_data = false + pre_integration_hook! = empty_hook + integration_hook = empty_hook + post_integration_hook = empty_hook end simulate( @@ -113,12 +120,16 @@ function run_sim( sim_consts.δt, T0, T, - n_steps_before_verlet_list_update, + sim.consts.n_steps_before_verlet_list_update, n_steps_before_snapshot, n_bundles, dir, save_data, + rl_params, + pre_integration_hook!, + integration_hook, + post_integration_hook, ) - return dir + return nothing end \ No newline at end of file diff --git a/src/setup.jl b/src/setup.jl index 376b064..f978f09 100644 --- a/src/setup.jl +++ b/src/setup.jl @@ -1,6 +1,12 @@ using Distributions: Uniform using Dates: now -using JSON3: JSON3 + +const DEFAULT_PACKING_RATIO = 0.5 +const DEFAULT_δt = 1e-5 +const DEFAULT_SKIN_TO_INTERACTION_R_RATIO = 3.0 +const DEFAULT_EXPORTS_DIR = "exports" +const DEFAULT_PARENT_DIR = "" +const DEFAULT_COMMENT = "" function initial_particle_grid_pos( i::Int64, j::Int64, grid_box_width::Float64, half_box_len::Float64 @@ -9,7 +15,7 @@ function initial_particle_grid_pos( return SVector(i * grid_box_width + term, j * grid_box_width + term) end -function generate_particles(grid_n::Int64, grid_box_width::Float64, half_box_len::Float64) +function gen_particles(grid_n::Int64, grid_box_width::Float64, half_box_len::Float64) particles = Vector{Particle}(undef, grid_n^2) id = 1 @@ -29,7 +35,7 @@ function generate_particles(grid_n::Int64, grid_box_width::Float64, half_box_len return particles end -function generate_particles(bundle::Bundle, n_particles::Int64) +function gen_particles(bundle::Bundle, n_particles::Int64) particles = Vector{Particle}(undef, n_particles) @simd for id in 1:n_particles @@ -39,27 +45,34 @@ function generate_particles(bundle::Bundle, n_particles::Int64) return particles end -function init_sim(; +function gen_sim_consts( n_particles::Int64, - v::Float64, - δt::Float64=1e-5, - packing_ratio::Float64=0.5, - comment::String="", - parent_dir::String="", + v₀::Float64, + δt::Float64=DEFAULT_δt, + packing_ratio::Float64=DEFAULT_PACKING_RATIO, + skin_to_interaction_r_ratio::Float64=DEFAULT_SKIN_TO_INTERACTION_R_RATIO, ) @assert n_particles > 0 - @assert v >= 0 + @assert v₀ >= 0 @assert δt in 1e-7:1e-7:1e-4 @assert packing_ratio > 0 μ = 1.0 D₀ = 1.0 - particle_diameter = 1.0 - Dᵣ = 3 * D₀ / (particle_diameter^2) + particle_radius = 1.0 + Dᵣ = 3 * D₀ / ((2 * particle_radius)^2) σ = 1.0 ϵ = 100.0 interaction_r = 2^(1 / 6) * σ + buffer = 6.0 + max_approach_after_one_integration_step = buffer * (2 * v₀ * δt) / interaction_r + @assert skin_to_interaction_r_ratio >= 1 + max_approach_after_one_integration_step + + skin_r = skin_to_interaction_r_ratio * interaction_r + n_steps_before_verlet_list_update = + (skin_to_interaction_r_ratio - 1) / max_approach_after_one_integration_step + grid_n = round(Int64, ceil(sqrt(n_particles))) n_particles = grid_n^2 @@ -67,47 +80,80 @@ function init_sim(; half_box_len = sqrt(n_particles * π / packing_ratio) * σ / 4 grid_box_width = 2 * half_box_len / grid_n - sim_consts = (; + return (; # Input n_particles, - v, + v₀, δt, packing_ratio, # Calculated μ, D₀, - particle_diameter, + particle_radius, Dᵣ, σ, ϵ, interaction_r, + skin_r, + n_steps_before_verlet_list_update, grid_n, half_box_len, grid_box_width, ) +end - particles = generate_particles(grid_n, grid_box_width, half_box_len) +function init_sim_with_sim_consts( + sim_consts, + exports_dir::String=DEFAULT_EXPORTS_DIR, + parent_dir::String=DEFAULT_PARENT_DIR, + comment::String=DEFAULT_COMMENT, +) + particles = gen_particles( + sim_consts.grid_n, sim_consts.grid_box_width, sim_consts.half_box_len + ) bundle = Bundle(n_particles, 1) save_snapshot!(bundle, 1, 0.0, particles) particles = nothing - dir = "exports" if length(parent_dir) > 0 - dir *= "/$parent_dir" + exports_dir *= "/$parent_dir" end - dir *= "/$(now())_N=$(n_particles)_v=$(v)_#$(rand(1000:9999))" + + start_datetime = now() + exports_dir *= "/$(start_datetime)_N=$(sim_consts.n_particles)_v=$(sim_consts.v₀)_#$(rand(1000:9999))" + if length(comment) > 0 - dir *= "_$comment" + exports_dir *= "_$comment" end - mkpath(dir) + mkpath(exports_dir) - open("$dir/sim_consts.json", "w") do f - JSON3.write(f, sim_consts) - end + task = @async write_struct_to_json(sim_consts, "$exports_dir/sim_consts") - save_bundle(dir, bundle, 1, 0.0) + save_bundle(exports_dir, bundle, 1, 0.0) - return dir + runs_dir = "$exports_dir/runs" + mkpath(runs_dir) + + wait(task) + + return exports_dir +end + +function init_sim(; + n_particles::Int64, + v₀::Float64, + δt::Float64=DEFAULT_δt, + packing_ratio::Float64=DEFAULT_PACKING_RATIO, + skin_to_interaction_r_ratio::Float64=DEFAULT_SKIN_TO_INTERACTION_R_RATIO, + exports_dir::String=DEFAULT_EXPORTS_DIR, + parent_dir::String=DEFAULT_PARENT_DIR, + comment::String=DEFAULT_COMMENT, +) + sim_consts = gen_sim_consts( + n_particles, v₀, δt, packing_ratio, skin_to_interaction_r_ratio + ) + + return init_sim_with_sim_consts(sim_consts, exports_dir, parent_dir, comment) end \ No newline at end of file diff --git a/src/simulation.jl b/src/simulation.jl index eb362e4..0afec22 100644 --- a/src/simulation.jl +++ b/src/simulation.jl @@ -40,7 +40,9 @@ function update_verlet_lists!(args, cl) return cl end -function euler!(args) +function euler!( + args, integration_hook::F, actions::Vector{Tuple{Float64,Float64}} +) where {F<:Function} for id1 in 1:(args.n_particles - 1) p1 = args.particles[id1] p1_c = p1.c @@ -66,21 +68,23 @@ function euler!(args) @simd for p in args.particles si, co = sincos(p.φ) p.tmp_c += SVector( - args.vδt * co + args.c₃ * rand_normal01(), - args.vδt * si + args.c₃ * rand_normal01(), + args.v₀δt * co + args.c₃ * rand_normal01(), + args.v₀δt * si + args.c₃ * rand_normal01(), ) p.φ += args.c₄ * rand_normal01() restrict_coordinates!(p, args.half_box_len) + integration_hook(p, actions) + p.c = p.tmp_c end return nothing end -wait(n::Nothing) = n +wait(::Nothing) = nothing function simulate( args, @@ -92,7 +96,11 @@ function simulate( n_bundles::Int64, dir::String, save_data::Bool, -) + rl_params::Union{ReCoRL.Params,Nothing}, + pre_integration_hook!::F, + integration_hook::F, + post_integration_hook::F, +) where {F<:Function} bundle_snapshot_counter = 0 task::Union{Task,Nothing} = nothing @@ -100,6 +108,8 @@ function simulate( cl = CellList(args.particles_c, args.box; parallel=false) cl = update_verlet_lists!(args, cl) + update_actions = true + start_time = now() println("Started simulation at $start_time.") @@ -124,7 +134,17 @@ function simulate( cl = update_verlet_lists!(args, cl) end - euler!(args) + update_actions = integration_step % rl_params.n_steps_before_actions_update == 0 + + if update_actions + pre_integration_hook!(rl_params) + end + + euler!(args, integration_hook, rl.params.actions) + + if update_actions + post_integration_hook(rl_params) + end end wait(task) diff --git a/src/utils.jl b/src/utils.jl deleted file mode 100644 index 16fcf30..0000000 --- a/src/utils.jl +++ /dev/null @@ -1,16 +0,0 @@ -using JSON3: JSON3 -using OrderedCollections: OrderedDict - -function struct_to_ordered_dict(s) - return OrderedDict(key => getfield(s, key) for key in propertynames(s)) -end - -function write_struct_to_json(s, path_without_extension::String) - ordered_dict = struct_to_ordered_dict(s) - - open("$path_without_extension.json", "w") do f - JSON3.write(f, ordered_dict) - end - - return nothing -end \ No newline at end of file