From 8a4bb5a9790cc4348174fb51187314d2f728eb3c Mon Sep 17 00:00:00 2001
From: Aman Goel <goelaman@amazon.com>
Date: Wed, 8 Nov 2023 11:00:14 -0800
Subject: [PATCH] [PCover] Update rl reward feedback

---
 .../runtime/statistics/CoverageStats.java     | 32 +++++++++----------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/statistics/CoverageStats.java b/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/statistics/CoverageStats.java
index ed971cfe9..daceb8d3c 100644
--- a/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/statistics/CoverageStats.java
+++ b/Src/PRuntimes/PSymRuntime/src/main/java/psym/runtime/statistics/CoverageStats.java
@@ -145,16 +145,16 @@ public void updateIterationCoverage(
     estimatedCoverage = estimatedCoverage.add(iterationCoverage);
     //        assert (estimatedCoverage.compareTo(BigDecimal.ONE) <= 0): "Error in path coverage
     // estimation";
-//    if (rewardMode != ChoiceLearningRewardMode.None) {
-//      for (int i = startDepth; i <= choiceDepth; i++) {
-//        CoverageChoiceDepthStats stats = perChoiceDepthStats.get(i);
-//        if (stats != null) {
-//          PSymGlobal.getChoiceLearningStats()
-//              .rewardIteration(
-//                  stats.getStateActions(), iterationCoverage.doubleValue(), rewardMode);
-//        }
-//      }
-//    }
+    if (rewardMode != ChoiceLearningRewardMode.None) {
+      for (int i = startDepth; i <= choiceDepth; i++) {
+        CoverageChoiceDepthStats stats = perChoiceDepthStats.get(i);
+        if (stats != null) {
+          PSymGlobal.getChoiceLearningStats()
+              .rewardIteration(
+                  stats.getStateActions(), iterationCoverage.doubleValue(), rewardMode);
+        }
+      }
+    }
   }
 
   /** Reset coverage statistics after a resumed run */
@@ -345,7 +345,7 @@ boolean isNotEmpty() {
   public static class CoverageChoiceDepthStats implements Serializable {
     BigDecimal pathCoverage;
     int numTotal;
-//    @Getter ChoiceQTable.ChoiceQTableKey stateActions;
+    @Getter ChoiceQTable.ChoiceQTableKey stateActions;
 
     CoverageChoiceDepthStats() {
       this(new BigDecimal(1), 0, new ChoiceQTable.ChoiceQTableKey());
@@ -361,7 +361,7 @@ private CoverageChoiceDepthStats(
         ChoiceQTable.ChoiceQTableKey inputStateActions) {
       this.pathCoverage = inputPathCoverage;
       this.numTotal = inputNumTotal;
-//      this.stateActions = inputStateActions;
+      this.stateActions = inputStateActions;
     }
 
     void update(
@@ -382,15 +382,15 @@ void update(
                 BigDecimal.valueOf(numExplored)
                     .divide(BigDecimal.valueOf(numTotal), 20, RoundingMode.FLOOR));
       }
-//      this.stateActions = chosenActions;
+      this.stateActions = chosenActions;
     }
 
     public void reset() {
       pathCoverage = new BigDecimal(1);
       numTotal = 0;
-//      if (stateActions != null) {
-//        stateActions.clear();
-//      }
+      if (stateActions != null) {
+        stateActions.clear();
+      }
     }
 
     public CoverageChoiceDepthStats getCopy() {