From dde6c3c08f42fc295476b34fcfa1b0c672e5d8ba Mon Sep 17 00:00:00 2001
From: Matthew Gerber <matthew.gerber@dexcom.com>
Date: Sat, 25 Nov 2023 19:28:49 -0500
Subject: [PATCH] Treat truncation as termination.

---
 run_configurations/CartPole train FA.run.xml | 2 +-
 src/rlai/core/environments/gymnasium.py      | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/run_configurations/CartPole train FA.run.xml b/run_configurations/CartPole train FA.run.xml
index 543fae50..4db780e2 100644
--- a/run_configurations/CartPole train FA.run.xml	
+++ b/run_configurations/CartPole train FA.run.xml	
@@ -13,7 +13,7 @@
     <option name="ADD_SOURCE_ROOTS" value="true" />
     <EXTENSION ID="PythonCoverageRunConfigurationExtension" runner="coverage.py" />
     <option name="SCRIPT_NAME" value="$PROJECT_DIR$/src/rlai/runners/trainer.py" />
-    <option name="PARAMETERS" value="--random-seed 12345 --agent rlai.gpi.state_action_value.ActionValueMdpAgent --gamma 0.95 --environment rlai.core.environments.gymnasium.Gym --gym-id CartPole-v1 --render-every-nth-episode 1000 --video-directory ~/Desktop/cartpole_videos --train-function rlai.gpi.temporal_difference.iteration.iterate_value_q_pi --mode SARSA --num-improvements 30000 --num-episodes-per-improvement 1 --num-updates-per-improvement 1 --epsilon 0.2 --q-S-A rlai.gpi.state_action_value.function_approximation.ApproximateStateActionValueEstimator --plot-model --plot-model-per-improvements 1000 --function-approximation-model rlai.gpi.state_action_value.function_approximation.models.sklearn.SKLearnSGD --loss squared_error --sgd-alpha 0.0 --learning-rate constant --eta0 0.00001 --feature-extractor rlai.core.environments.gymnasium.CartpoleFeatureExtractor --make-final-policy-greedy True --num-improvements-per-plot 1000 --num-improvements-per-checkpoint 1000 --checkpoint-path ~/Desktop/cartpole_checkpoint/cartpole_checkpoint.pickle --save-agent-path ~/Desktop/cartpole_agent.pickle --log INFO" />
+    <option name="PARAMETERS" value="--random-seed 12345 --agent rlai.gpi.state_action_value.ActionValueMdpAgent --gamma 0.95 --environment rlai.core.environments.gymnasium.Gym --gym-id CartPole-v1 --render-every-nth-episode 1000 --video-directory ~/Desktop/cartpole_videos --train-function rlai.gpi.temporal_difference.iteration.iterate_value_q_pi --mode SARSA --num-improvements 50000 --num-episodes-per-improvement 1 --num-updates-per-improvement 1 --epsilon 0.2 --q-S-A rlai.gpi.state_action_value.function_approximation.ApproximateStateActionValueEstimator --plot-model --plot-model-per-improvements 1000 --function-approximation-model rlai.gpi.state_action_value.function_approximation.models.sklearn.SKLearnSGD --loss squared_error --sgd-alpha 0.0 --learning-rate constant --eta0 0.00001 --feature-extractor rlai.core.environments.gymnasium.CartpoleFeatureExtractor --make-final-policy-greedy True --num-improvements-per-plot 1000 --num-improvements-per-checkpoint 1000 --checkpoint-path ~/Desktop/cartpole_checkpoint/cartpole_checkpoint.pickle --save-agent-path ~/Desktop/cartpole_agent.pickle --log INFO" />
     <option name="SHOW_COMMAND_LINE" value="false" />
     <option name="EMULATE_TERMINAL" value="false" />
     <option name="MODULE_MODE" value="false" />
diff --git a/src/rlai/core/environments/gymnasium.py b/src/rlai/core/environments/gymnasium.py
index bb5c1a9c..5efc24f9 100644
--- a/src/rlai/core/environments/gymnasium.py
+++ b/src/rlai/core/environments/gymnasium.py
@@ -1,3 +1,4 @@
+import logging
 import math
 import os
 import warnings
@@ -233,7 +234,12 @@ def advance(
             else:
                 fuel_used = required_fuel
 
-        observation, reward, terminated, _, _ = self.gym_native.step(action=gym_action)
+        observation, reward, terminated, truncated, _ = self.gym_native.step(action=gym_action)
+
+        # truncation is a special case of termination
+        if truncated and not terminated:
+            logging.info(f'Episode was truncated after {t + 1} step(s). Terminating.')
+            terminated = truncated
 
         # update fuel remaining if needed
         fuel_remaining = None