Envs cleanup

P6-Pool · Apr 10, 2024 · 8fe5001 · 8fe5001
1 parent f482b0e
commit 8fe5001
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 192 deletions.
diff --git a/src/fastfiz_env/__init__.py b/src/fastfiz_env/__init__.py
@@ -37,17 +37,26 @@
 register(
     id="BaseFastFiz-v0",
     entry_point="fastfiz_env.envs:BaseFastFiz",
+    additional_wrappers=(
+        utils.wrappers.MaxEpisodeStepsInjectionWrapper.wrapper_spec(),
+    ),
 )
 
 register(
     id="BaseRLFastFiz-v0",
     entry_point="fastfiz_env.envs:BaseRLFastFiz",
+    additional_wrappers=(
+        utils.wrappers.MaxEpisodeStepsInjectionWrapper.wrapper_spec(),
+    ),
 )
 
 
 register(
     id="VelocityFastFiz-v0",
     entry_point="fastfiz_env.envs:VelocityFastFiz",
+    additional_wrappers=(
+        utils.wrappers.MaxEpisodeStepsInjectionWrapper.wrapper_spec(),
+    ),
 )
 
 register(

diff --git a/src/fastfiz_env/envs/event_fastfiz.py b/src/fastfiz_env/envs/event_fastfiz.py
diff --git a/src/fastfiz_env/utils/reward_functions/common/exponential_velocity_reward.py b/src/fastfiz_env/utils/reward_functions/common/exponential_velocity_reward.py
@@ -17,5 +17,5 @@ def reward(
         """
         Reward function that gives a reward based on velocity of the action.
         """
-        reward = np.interp(action[4], [0, 10], [0, 1])
-        return reward**10
+        reward = float(np.interp(action[4], [0, 10], [0, 1]))
+        return reward**2