diff --git a/agents/C_trajectory_following.py b/agents/C_trajectory_following.py
index 05362541..0361f326 100755
--- a/agents/C_trajectory_following.py
+++ b/agents/C_trajectory_following.py
@@ -35,7 +35,7 @@
 ###### ENVIRONMENT PARAMS ##########
 TIMESTEPS = 2.5e6
 N_ENVS = 20
-EPISODE_LEN_SEC = 30
+EPISODE_LEN_SEC = 10
 ####################################
 
 ###### HYPERPARAMS #################
@@ -44,7 +44,7 @@
 K_WP = 20
 K_S = 0.05
 MAX_REWARD_DISTANCE = 0.03
-WAYPOINT_DIST_TOL = 0.3
+WAYPOINT_DIST_TOL = 0.12
 ####################################
 
 
@@ -87,7 +87,7 @@ def run(output_folder=OUTPUT_FOLDER,
     t_traj, init_wp = init_targets()
 
     # random number in range 10-99
-    if train:
+    if train or test:
         output_folder = f"{output_folder}/k_p={k_p}_k_wp={k_wp}_k_s={k_s}_max_reward_distance={max_reward_distance}_waypoint_dist_tol={waypoint_dist_tol}"
     print(f"Output folder: {output_folder}")