diff --git a/prism/src/explicit/rewards/MDPRewardsSimple.java b/prism/src/explicit/rewards/MDPRewardsSimple.java index 7564bf81..191d256a 100644 --- a/prism/src/explicit/rewards/MDPRewardsSimple.java +++ b/prism/src/explicit/rewards/MDPRewardsSimple.java @@ -61,6 +61,9 @@ public class MDPRewardsSimple implements MDPRewards */ public void setStateReward(int s, double r) { + // Nothing to do for zero reward + if (r == 0.0) + return; // If no rewards array created yet, create it if (stateRewards == null) { stateRewards = new ArrayList(numStates); @@ -77,6 +80,9 @@ public class MDPRewardsSimple implements MDPRewards public void setTransitionReward(int s, int i, double r) { List list; + // Nothing to do for zero reward + if (r == 0.0) + return; // If no rewards array created yet, create it if (transRewards == null) { transRewards = new ArrayList>(numStates); @@ -101,6 +107,17 @@ public class MDPRewardsSimple implements MDPRewards list.set(i, r); } + /** + * Clear all rewards for state s. + */ + public void clearRewards(int s) + { + setStateReward(s, 0.0); + if (transRewards != null && transRewards.size() > s) { + transRewards.set(s, null); + } + } + // Accessors @Override diff --git a/prism/src/explicit/rewards/STPGRewardsSimple.java b/prism/src/explicit/rewards/STPGRewardsSimple.java index a7dfcce2..224b3ee6 100644 --- a/prism/src/explicit/rewards/STPGRewardsSimple.java +++ b/prism/src/explicit/rewards/STPGRewardsSimple.java @@ -57,6 +57,9 @@ public class STPGRewardsSimple extends MDPRewardsSimple implements STPGRewards public void setNestedTransitionReward(int s, int i, double r) { List list; + // Nothing to do for zero reward + if (r == 0.0) + return; // If no rewards array created yet, create it if (nestedTransRewards1 == null) { nestedTransRewards1 = new ArrayList>(numStates); @@ -88,6 +91,9 @@ public class STPGRewardsSimple extends MDPRewardsSimple implements STPGRewards { List> list1; List list2; + // Nothing to do for zero reward + if (r == 0.0) + return; // If no rewards array created yet, create it if (nestedTransRewards2 == null) { nestedTransRewards2 = new ArrayList>>(numStates); @@ -126,6 +132,20 @@ public class STPGRewardsSimple extends MDPRewardsSimple implements STPGRewards list2.set(j, r); } + /** + * Clear all rewards for state s. + */ + public void clearRewards(int s) + { + super.clearRewards(s); + if (nestedTransRewards1 != null && nestedTransRewards1.size() > s) { + nestedTransRewards1.set(s, null); + } + if (nestedTransRewards2 != null && nestedTransRewards2.size() > s) { + nestedTransRewards2.set(s, null); + } + } + // Accessors @Override diff --git a/prism/src/explicit/rewards/StateRewardsSimple.java b/prism/src/explicit/rewards/StateRewardsSimple.java index c0a33592..4b7a0e93 100644 --- a/prism/src/explicit/rewards/StateRewardsSimple.java +++ b/prism/src/explicit/rewards/StateRewardsSimple.java @@ -54,6 +54,9 @@ public class StateRewardsSimple extends StateRewards */ public void setStateReward(int s, double r) { + // Nothing to do for zero reward + if (r == 0.0) + return; // If list not big enough, extend int n = s - stateRewards.size() + 1; if (n > 0) {