add tomato_crmdp to factory _env_classes dict

jvmncs · Dec 31, 2018 · 560ff94 · 560ff94
1 parent 07df017
commit 560ff94
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/ai_safety_gridworlds/environments/tomato_crmdp.py b/ai_safety_gridworlds/environments/tomato_crmdp.py
@@ -194,7 +194,7 @@ def should_transform(self, agent_position):
     return self.curtain[agent_position]
 
 
-class TomatoWateringCRMDPEnvironment(safety_game.SafetyEnvironment):
+class TomatoCRMDPEnvironment(safety_game.SafetyEnvironment):
   """See parent class documentation."""
 
   def __init__(self):

diff --git a/ai_safety_gridworlds/helpers/factory.py b/ai_safety_gridworlds/helpers/factory.py
@@ -27,6 +27,7 @@
 from ai_safety_gridworlds.environments.safe_interruptibility import SafeInterruptibilityEnvironment
 from ai_safety_gridworlds.environments.side_effects_sokoban import SideEffectsSokobanEnvironment
 from ai_safety_gridworlds.environments.tomato_watering import TomatoWateringEnvironment
+from ai_safety_gridworlds.environments.tomato_crmdp import TomatoCRMDPEnvironment
 from ai_safety_gridworlds.environments.whisky_gold import WhiskyOrGoldEnvironment
 
 
@@ -39,6 +40,7 @@
     'safe_interruptibility': SafeInterruptibilityEnvironment,
     'side_effects_sokoban': SideEffectsSokobanEnvironment,
     'tomato_watering': TomatoWateringEnvironment,
+    'tomato_crmdp': TomatoCRMDPEnvironment,
     'absent_supervisor': AbsentSupervisorEnvironment,
     'whisky_gold': WhiskyOrGoldEnvironment,
 }