diff --git a/src/RL/RL.jl b/src/RL/RL.jl
index 5da861c..9a3ab5f 100644
--- a/src/RL/RL.jl
+++ b/src/RL/RL.jl
@@ -1,6 +1,6 @@
 module RL
 
-export run_rl, LocalCOMWithAdditionalShapeRewardEnv, OriginEnv
+export run_rl, LocalCOMWithAdditionalShapeRewardEnv, OriginEnv, NearestNeighbourEnv
 
 using Base: OneTo
 
@@ -29,8 +29,8 @@ include("Reward.jl")
 function gen_agent(
     n_states::Int64, n_actions::Int64, ϵ_stable::Float64, reward_discount::Float64
 )
-    # TODO: Optimize warmup and decay
-    warmup_steps = 500_000
+    # TODO: Optimize warming up and decay
+    warmup_steps = 400_000
     decay_steps = 5_000_000
 
     policy = QBasedPolicy(;
@@ -188,5 +188,6 @@ end
 
 include("Envs/LocalCOMWithAdditionalShapeRewardEnv.jl")
 include("Envs/OriginEnv.jl")
+include("Envs/NearestNeighbourEnv.jl")
 
 end # module
\ No newline at end of file
diff --git a/src/ReCo.jl b/src/ReCo.jl
index d4e0072..8c44e7f 100644
--- a/src/ReCo.jl
+++ b/src/ReCo.jl
@@ -1,7 +1,13 @@
 module ReCo
 
 export init_sim,
-    run_sim, run_rl, animate, plot_snapshot, LocalCOMWithAdditionalShapeRewardEnv, OriginEnv
+    run_sim,
+    run_rl,
+    animate,
+    plot_snapshot,
+    LocalCOMWithAdditionalShapeRewardEnv,
+    OriginEnv,
+    NearestNeighbourEnv
 
 using StaticArrays: SVector
 using JLD2: JLD2