From fe7c65c4a3e55595ba11543dffc0bd233de06768 Mon Sep 17 00:00:00 2001 From: Mo8it Date: Fri, 14 Jan 2022 13:01:14 +0100 Subject: [PATCH] Fixed elliptical_distance --- src/RL/EnvHelper.jl | 2 +- src/RL/Hooks.jl | 10 ++++++---- src/RL/LocalCOMEnv.jl | 18 ++++++++++++++---- src/RL/RL.jl | 11 ++--------- src/Shape.jl | 36 ++++++++++++++++++++++++------------ 5 files changed, 47 insertions(+), 30 deletions(-) diff --git a/src/RL/EnvHelper.jl b/src/RL/EnvHelper.jl index 9eb8d39..bcb36f5 100644 --- a/src/RL/EnvHelper.jl +++ b/src/RL/EnvHelper.jl @@ -41,7 +41,7 @@ struct EnvHelperSharedProps{H<:AbstractHook} end function gen_env_helper(::Env, env_helper_params::EnvHelperSharedProps; args) - return method_not_implemented() + return ReCo.method_not_implemented() end function get_env_agent_hook(env_helper::EnvHelper) diff --git a/src/RL/Hooks.jl b/src/RL/Hooks.jl index a5311d7..edd9ae5 100644 --- a/src/RL/Hooks.jl +++ b/src/RL/Hooks.jl @@ -1,19 +1,21 @@ +using ..ReCo: Particle + function pre_integration_hook(::EnvHelper) - return method_not_implemented() + return ReCo.method_not_implemented() end function state_update_helper_hook( ::EnvHelper, id1::Int64, id2::Int64, r⃗₁₂::SVector{2,Float64} ) - return method_not_implemented() + return ReCo.method_not_implemented() end function state_update_hook(::EnvHelper, particles::Vector{Particle}) - return method_not_implemented() + return ReCo.method_not_implemented() end function update_reward!(::Env, ::EnvHelper, particle::Particle) - return method_not_implemented() + return ReCo.method_not_implemented() end function update_table_and_actions_hook( diff --git a/src/RL/LocalCOMEnv.jl b/src/RL/LocalCOMEnv.jl index 80b8e5b..8c2fef3 100644 --- a/src/RL/LocalCOMEnv.jl +++ b/src/RL/LocalCOMEnv.jl @@ -1,5 +1,7 @@ export LocalCOMEnv +using ..ReCo: Particle + struct LocalCOMEnv <: Env shared::EnvSharedProps @@ -50,6 +52,7 @@ mutable struct LocalCOMEnvHelper <: EnvHelper add_shape_reward_term::Bool + center_of_mass::SVector{2,Float64} gyration_tensor_eigvec_to_smaller_eigval::SVector{2,Float64} gyration_tensor_eigvec_to_bigger_eigval::SVector{2,Float64} @@ -71,6 +74,7 @@ mutable struct LocalCOMEnvHelper <: EnvHelper false, SVector(0.0, 0.0), SVector(0.0, 0.0), + SVector(0.0, 0.0), half_box_len, max_elliptical_distance, ) @@ -125,7 +129,7 @@ function state_update_hook(env_helper::LocalCOMEnvHelper, particles::Vector{Part vec_to_local_center_of_mass = env_helper.vec_to_neighbour_sums[id] / n_neighbours - distance = norm2d(vec_to_local_center_of_mass) + distance = ReCo.norm2d(vec_to_local_center_of_mass) env_helper.distances_to_local_center_of_mass[id] = distance @@ -135,7 +139,7 @@ function state_update_hook(env_helper::LocalCOMEnvHelper, particles::Vector{Part si, co = sincos(particles[id].φ) - direction_angle = angle2(SVector(co, si), vec_to_local_center_of_mass) + direction_angle = ReCo.angle2(SVector(co, si), vec_to_local_center_of_mass) direction_angle_state = find_state_interval( direction_angle, env.direction_angle_state_space @@ -158,7 +162,11 @@ function state_update_hook(env_helper::LocalCOMEnvHelper, particles::Vector{Part #println(mean_distance_to_local_center_of_mass / env_helper.max_distance_to_local_center_of_mass) # TODO: Remove end - v1, v2 = Shape.gyration_tensor_eigvecs(particles, env_helper.half_box_len) + env_helper.center_of_mass = ReCo.center_of_mass(particles, env_helper.half_box_len) + + v1, v2 = ReCo.gyration_tensor_eigvecs( + particles, env_helper.half_box_len, env_helper.center_of_mass + ) env_helper.gyration_tensor_eigvec_to_smaller_eigval = v1 env_helper.gyration_tensor_eigvec_to_bigger_eigval = v2 @@ -191,10 +199,12 @@ function update_reward!(env::LocalCOMEnv, env_helper::LocalCOMEnvHelper, particl if env_helper.add_shape_reward_term elliptical_distance = ReCo.elliptical_distance( - particle, + particle.c, + env_helper.center_of_mass, env_helper.gyration_tensor_eigvec_to_smaller_eigval, env_helper.gyration_tensor_eigvec_to_bigger_eigval, env_helper.shared.goal_gyration_tensor_eigvals_ratio, + env_helper.half_box_len, ) reward += unnormalized_reward( diff --git a/src/RL/RL.jl b/src/RL/RL.jl index 5287f8e..7f185b6 100644 --- a/src/RL/RL.jl +++ b/src/RL/RL.jl @@ -12,14 +12,7 @@ using LoopVectorization: @turbo using Random: Random using ProgressMeter: @showprogress -using ..ReCo: - ReCo, - Particle, - angle2, - norm2d, - Shape, - DEFAULT_SKIN_TO_INTERACTION_R_RATIO, - method_not_implemented +using ..ReCo: ReCo const INITIAL_STATE_IND = 1 const INITIAL_REWARD = 0.0 @@ -67,7 +60,7 @@ function run_rl(; n_particles::Int64=100, seed::Int64=42, ϵ_stable::Float64=0.0001, - skin_to_interaction_r_ratio::Float64=DEFAULT_SKIN_TO_INTERACTION_R_RATIO, + skin_to_interaction_r_ratio::Float64=ReCo.DEFAULT_SKIN_TO_INTERACTION_R_RATIO, packing_ratio=0.22, ) where {E<:Env} @assert 0.0 <= goal_gyration_tensor_eigvals_ratio <= 1.0 diff --git a/src/Shape.jl b/src/Shape.jl index 854fe8c..9975b5c 100644 --- a/src/Shape.jl +++ b/src/Shape.jl @@ -6,7 +6,7 @@ export center_of_mass, using StaticArrays: SVector, SMatrix using LinearAlgebra: eigvals, eigvecs, Hermitian, dot -using ..ReCo: Particle, restrict_coordinate, restrict_coordinates +using ..ReCo: Particle function project_to_unit_circle(x::Float64, half_box_len::Float64) φ = (x + half_box_len) * π / half_box_len @@ -18,7 +18,7 @@ end function project_back_from_unit_circle(θ::T, half_box_len::Float64) where {T<:Real} x = θ * half_box_len / π - half_box_len - return restrict_coordinate(x, half_box_len) + return ReCo.restrict_coordinate(x, half_box_len) end function center_of_mass_from_proj_sums( @@ -71,15 +71,15 @@ function center_of_mass(particles::Vector{Particle}, half_box_len::Float64) return center_of_mass_from_proj_sums(x_proj_sum, y_proj_sum, half_box_len) end -function gyration_tensor(particles::Vector{Particle}, half_box_len::Float64) - COM = center_of_mass(particles, half_box_len) - +function gyration_tensor( + particles::Vector{Particle}, half_box_len::Float64, COM::SVector{2,Float64} +) S11 = 0.0 S12 = 0.0 S22 = 0.0 for p in particles - shifted_c = restrict_coordinates(p.c - COM, half_box_len) + shifted_c = ReCo.restrict_coordinates(p.c - COM, half_box_len) S11 += shifted_c[1]^2 S12 += shifted_c[1] * shifted_c[2] @@ -89,14 +89,22 @@ function gyration_tensor(particles::Vector{Particle}, half_box_len::Float64) return Hermitian(SMatrix{2,2}(S11, S12, S12, S22)) end +function gyration_tensor(particles::Vector{Particle}, half_box_len::Float64) + COM = center_of_mass(particles, half_box_len) + + return gyration_tensor(particles, half_box_len, COM) +end + function gyration_tensor_eigvals_ratio(particles::Vector{Particle}, half_box_len::Float64) g_tensor = gyration_tensor(particles, half_box_len) ev = eigvals(g_tensor) # Eigenvalues are sorted return ev[1] / ev[2] end -function gyration_tensor_eigvecs(particles::Vector{Particle}, half_box_len::Float64) - g_tensor = gyration_tensor(particles, half_box_len) +function gyration_tensor_eigvecs( + particles::Vector{Particle}, half_box_len::Float64, COM::SVector{2,Float64} +) + g_tensor = gyration_tensor(particles, half_box_len, COM) eig_vecs = eigvecs(g_tensor) v1 = eig_vecs[:, 1] @@ -106,15 +114,19 @@ function gyration_tensor_eigvecs(particles::Vector{Particle}, half_box_len::Floa end function elliptical_distance( - particle::Particle, + v::SVector{2,Float64}, + COM::SVector{2,Float64}, gyration_tensor_eigvec_to_smaller_eigval::SVector{2,Float64}, gyration_tensor_eigvec_to_bigger_eigval::SVector{2,Float64}, goal_gyration_tensor_eigvals_ratio::Float64, + half_box_len::Float64, ) - cx′ = dot(particle.c, gyration_tensor_eigvec_to_bigger_eigval) - cy′ = dot(particle.c, gyration_tensor_eigvec_to_smaller_eigval) + v′ = ReCo.minimum_image(v - COM, half_box_len) - return cx′^2 + (cy′ / goal_gyration_tensor_eigvals_ratio)^2 + x = dot(v′, gyration_tensor_eigvec_to_bigger_eigval) + y = dot(v′, gyration_tensor_eigvec_to_smaller_eigval) + + return x^2 + (y / goal_gyration_tensor_eigvals_ratio)^2 end end # module \ No newline at end of file