Repository: pyrddlgym-project/pyRDDLGym (language: python)

Issue:
Jax front end overhaul
Major overhaul of the front end api for agents and Jax planner.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index da693471..f60cc2a0 100644
--- a/README.md
+++ b/README.md
@@ -88,7 +88,7 @@ pip install pyRDDLGym
 
 #### Known issues
 There are two known issues not documented with RDDL
-1. The minus (-) arithmatic operation must have spaces on both sides,
+1. The minus (-) arithmetic operation must have spaces on both sides,
 otherwise there is ambiguity is whether it is a mathematical operation of a fluent name.
 2. Aggregation union precedence requires for encapsulating parentheses, e.g., (sum_{}[]).
 
@@ -99,7 +99,7 @@ In addition, we supply two simple agents, to illustrate interaction with the env
 ```python
 from pyRDDLGym import RDDLEnv
 from pyRDDLGym import ExampleManager
-from pyRDDLGym.Policies.Agents import RandomAgent
+from pyRDDLGym.Core.Policies.Agents import RandomAgent
 ```
 The list of examples can be obtained through the ExampleManager object:
 ```python
diff --git a/docs/start.rst b/docs/start.rst
index 7a1dd7ad..0cb5f176 100644
--- a/docs/start.rst
+++ b/docs/start.rst
@@ -78,7 +78,7 @@ If the ``env.render()`` function is used, we will also see a window pop up rende
 
     from pyRDDLGym import RDDLEnv
     from pyRDDLGym import ExampleManager
-    from pyRDDLGym.Policies.Agents import RandomAgent
+    from pyRDDLGym.Core.Policies.Agents import RandomAgent
 
     # get the environment info
     EnvInfo = ExampleManager.GetEnvInfo('MarsRover')
diff --git a/pyRDDLGym/Examples/Traffic/Experiments/eval.py b/pyRDDLGym/Examples/Traffic/Experiments/eval.py
index e2b1a73b..4474351b 100644
--- a/pyRDDLGym/Examples/Traffic/Experiments/eval.py
+++ b/pyRDDLGym/Examples/Traffic/Experiments/eval.py
@@ -4,11 +4,12 @@
 
 from pyRDDLGym import RDDLEnv
 from pyRDDLGym import ExampleManager
-from pyRDDLGym.Policies.Agents import RandomAgent
+from pyRDDLGym.Core.Policies.Agents import RandomAgent
 
 from learned_plans.two_cycle import plan as JAX2CyclePlan
 
-parser = argparse.ArgumentParser(description='Evaluating benchmarks for a single-intersection traffic scenario')
+parser = argparse.ArgumentParser(
+    description='Evaluating benchmarks for a single-intersection traffic scenario')
 parser.add_argument('agent', type=str, choices=[
     'webster',
     'random',
@@ -16,7 +17,8 @@
     'max',
     'jax2cycle'
 ])
-parser.add_argument('-r', '--render', action='store_true', help='Whether to render (visualize) the rollout')
+parser.add_argument('-r', '--render', action='store_true',
+                    help='Whether to render (visualize) the rollout')
 parser.add_argument('-w', '--warmup', type=int, default=-1)
 args = parser.parse_args()
 
@@ -38,29 +40,42 @@
 # set up agent
 webster_change_ts = {32, 36, 69, 73, 106, 110, 143, 147}
 webster_CT = 148
+
+
 def get_webster_action(step):
-    return {'advance___i0': step%webster_CT in webster_change_ts}
+    return {'advance___i0': step % webster_CT in webster_change_ts}
+
 
 random_agent = RandomAgent(action_space=myEnv.action_space,
                            num_actions=myEnv.numConcurrentActions)
 flips = bernoulli.rvs(0.5, size=myEnv.horizon)
+
+
 def get_random_action(step):
     return {'advance___i0': flips[step]}
     return random_agent.sample_action()
 
+
 min_change_ts = {5, 9, 15, 19, 25, 29, 35, 39}
 min_CT = 40
+
+
 def get_min_action(step):
-    return {'advance___i0': step%min_CT in min_change_ts}
+    return {'advance___i0': step % min_CT in min_change_ts}
+
 
 max_change_ts = {59, 63, 123, 127, 187, 191, 251, 255}
 max_CT = 256
+
+
 def get_max_action(step):
-    return {'advance___i0': step%min_CT in max_change_ts}
+    return {'advance___i0': step % min_CT in max_change_ts}
+
 
 def get_jax_2_cycle_action(step):
     return {'advance___i0': JAX2CyclePlan[step]}
 
+
 if args.agent == 'random':
     get_action = get_random_action
 elif args.agent == 'webster':
@@ -92,7 +107,8 @@ def get_jax_2_cycle_action(step):
 
 # evaluate
 for step in range(myEnv.horizon-args.warmup):
-    if args.render: myEnv.render()
+    if args.render:
+        myEnv.render()
     action = get_action(step)
     next_state, reward, done, info = myEnv.step(action)
     total_reward += reward


Test patch:
diff --git a/pyRDDLGym/tests/seeding/Seedtest.py b/pyRDDLGym/tests/seeding/Seedtest.py
index 306c54f9..6cd1e9ed 100644
--- a/pyRDDLGym/tests/seeding/Seedtest.py
+++ b/pyRDDLGym/tests/seeding/Seedtest.py
@@ -2,7 +2,7 @@
 
 from pyRDDLGym import RDDLEnv
 from pyRDDLGym.Core.Env.RDDLEnvSeeder import RDDLEnvSeeder
-from pyRDDLGym.Policies.Agents import NoOpAgent as RDDLAgent
+from pyRDDLGym.Core.Policies.Agents import NoOpAgent as RDDLAgent
 
 
 def Test_instance(dom_path, inst_path, episodes=1):
@@ -12,7 +12,7 @@ def Test_instance(dom_path, inst_path, episodes=1):
                             debug=False,
                             seeds=[42, 43])
     agent = RDDLAgent(action_space=myEnv.action_space,
-                        num_actions=myEnv.numConcurrentActions)
+                      num_actions=myEnv.numConcurrentActions)
 
     for episode in range(episodes):
         total_reward = 0
@@ -48,7 +48,6 @@ def main(path):
     Test_instance(dom_path, inst_path, 3)
 
 
-
 if __name__ == "__main__":
 
     path = os.path.dirname(os.path.abspath(__file__))
diff --git a/pyRDDLGym/tests/testGymDrone.py b/pyRDDLGym/tests/testGymDrone.py
index ecf2a338..5f7f943e 100644
--- a/pyRDDLGym/tests/testGymDrone.py
+++ b/pyRDDLGym/tests/testGymDrone.py
@@ -1,13 +1,15 @@
 from pyRDDLGym.Core.Env import RDDLEnv as RDDLEnv
-from pyRDDLGym.Policies.Agents import RandomAgent
+from pyRDDLGym.Core.Policies.Agents import RandomAgent
 
 FOLDER = './Competition/Drone/Drone_con/'
 
 
 def main():
     steps = 1000
-    myEnv = RDDLEnv.RDDLEnv(domain=FOLDER + 'domain.rddl', instance=FOLDER + 'insta0.rddl', is_grounded=False)
-    agent = RandomAgent(action_space=myEnv.action_space, num_actions=myEnv.NumConcurrentActions)
+    myEnv = RDDLEnv.RDDLEnv(domain=FOLDER + 'domain.rddl',
+                            instance=FOLDER + 'insta0.rddl')
+    agent = RandomAgent(action_space=myEnv.action_space,
+                        num_actions=myEnv.numConcurrentActions)
 
     from pprint import pprint
     pprint(vars(myEnv.model))
diff --git a/pyRDDLGym/tests/test_basic.py b/pyRDDLGym/tests/test_basic.py
index 79e6f505..a18d98c5 100644
--- a/pyRDDLGym/tests/test_basic.py
+++ b/pyRDDLGym/tests/test_basic.py
@@ -1,16 +1,18 @@
 from pyRDDLGym.Core.Env import RDDLEnv as RDDLEnv
-from pyRDDLGym.Policies.Agents import RandomAgent
+from pyRDDLGym.Core.Policies.Agents import RandomAgent
 
-FOLDER = 'Competition/MountainCar/'
+FOLDER = 'pyRDDLGym/Examples/MountainCar/'
 
 
 def main():
-    myEnv = RDDLEnv.RDDLEnv(domain=FOLDER + 'domain.rddl', instance=FOLDER + 'instance0.rddl', is_grounded=False)
-    agent = RandomAgent(action_space=myEnv.action_space, num_actions=myEnv.NumConcurrentActions)
-    
+    myEnv = RDDLEnv.RDDLEnv(domain=FOLDER + 'domain.rddl',
+                            instance=FOLDER + 'instance0.rddl')
+    agent = RandomAgent(action_space=myEnv.action_space,
+                        num_actions=myEnv.numConcurrentActions)
+
     from pprint import pprint
     pprint(vars(myEnv.model))
-    
+
     total_reward = 0
     state = myEnv.reset()
     for step in range(myEnv.model.horizon):
diff --git a/pyRDDLGym/tests/unit_test_mars_rover.py b/pyRDDLGym/tests/unit_test_mars_rover.py
index 2469cb91..b8b73b63 100644
--- a/pyRDDLGym/tests/unit_test_mars_rover.py
+++ b/pyRDDLGym/tests/unit_test_mars_rover.py
@@ -1,8 +1,8 @@
 from pyRDDLGym.Core.Env import RDDLEnv as RDDLEnv
-from pyRDDLGym.Policies.Agents import RandomAgent
+from pyRDDLGym.Core.Policies.Agents import RandomAgent
 import math
 
-FOLDER = 'Competition/Mars_rover/'
+FOLDER = 'pyRDDLGym/Examples/MarsRover/'
 
 
 def mars_rover_update(state, action):
@@ -24,14 +24,14 @@ def mars_rover_update(state, action):
         u_y_d2 = 0.05 * action['power-y_d2'] / power_d2
     else:
         u_y_d2 = action['power-y_d2']
-        
+
     next_interm = {'power_d1': power_d1,
-                    'power_d2': power_d2,
-                    'u-x_d1': u_x_d1,
-                    'u-y_d1': u_y_d1,
-                    'u-x_d2': u_x_d2,
-                    'u-y_d2': u_y_d2}
-    
+                   'power_d2': power_d2,
+                   'u-x_d1': u_x_d1,
+                   'u-y_d1': u_y_d1,
+                   'u-x_d2': u_x_d2,
+                   'u-y_d2': u_y_d2}
+
     new_vel_x_d1 = state['vel-x_d1'] + 0.1 * u_x_d1
     new_vel_y_d1 = state['vel-y_d1'] + 0.1 * u_y_d1
     new_vel_x_d2 = state['vel-x_d2'] + 0.1 * u_x_d2
@@ -40,14 +40,16 @@ def mars_rover_update(state, action):
     new_pos_y_d1 = state['pos-y_d1'] + 0.1 * state['vel-y_d1']
     new_pos_x_d2 = state['pos-x_d2'] + 0.1 * state['vel-x_d2']
     new_pos_y_d2 = state['pos-y_d2'] + 0.1 * state['vel-y_d2']
-    
+
     print(action)
     new_mineral_harvested_m1 = state['mineral-harvested_m1'] or (
-        (math.sqrt((state['pos-x_d1'] - 5) ** 2 + (state['pos-y_d1'] - 5) ** 2) < 6 and action['harvest_d1'])
+        (math.sqrt((state['pos-x_d1'] - 5) ** 2 +
+         (state['pos-y_d1'] - 5) ** 2) < 6 and action['harvest_d1'])
         or (math.sqrt((state['pos-x_d2'] - 5) ** 2 + (state['pos-y_d2'] - 5) ** 2) < 6 and action['harvest_d2'])
     )
     new_mineral_harvested_m2 = state['mineral-harvested_m2'] or (
-        (math.sqrt((state['pos-x_d1'] - (-8)) ** 2 + (state['pos-y_d1'] - (-8)) ** 2) < 8 and action['harvest_d1'])
+        (math.sqrt((state['pos-x_d1'] - (-8)) ** 2 +
+         (state['pos-y_d1'] - (-8)) ** 2) < 8 and action['harvest_d1'])
         or (math.sqrt((state['pos-x_d2'] - (-8)) ** 2 + (state['pos-y_d2'] - (-8)) ** 2) < 8 and action['harvest_d2'])
     )
 
@@ -61,36 +63,41 @@ def mars_rover_update(state, action):
                   'pos-y_d2': new_pos_y_d2,
                   'mineral-harvested_m1': new_mineral_harvested_m1,
                   'mineral-harvested_m2': new_mineral_harvested_m2}
-    
+
     mineral_m1_cond = (
         8 if (
-            (math.sqrt((state['pos-x_d1'] - 5) ** 2 + (state['pos-y_d1'] - 5) ** 2) < 6 and action['harvest_d1'])
+            (math.sqrt((state['pos-x_d1'] - 5) ** 2 +
+             (state['pos-y_d1'] - 5) ** 2) < 6 and action['harvest_d1'])
             or (math.sqrt((state['pos-x_d2'] - 5) ** 2 + (state['pos-y_d2'] - 5) ** 2) < 6 and action['harvest_d2'])
-        ) 
+        )
         else 0
     )
-    
+
     mineral_m2_cond = (
         5 if (
-            (math.sqrt((state['pos-x_d1'] - -8) ** 2 + (state['pos-y_d1'] - -8) ** 2) < 8 and action['harvest_d1'])
+            (math.sqrt((state['pos-x_d1'] - -8) ** 2 +
+             (state['pos-y_d1'] - -8) ** 2) < 8 and action['harvest_d1'])
             or (math.sqrt((state['pos-x_d2'] - -8) ** 2 + (state['pos-y_d2'] - -8) ** 2) < 8 and action['harvest_d2'])
-        ) 
+        )
         else 0
     )
-    
-    reward = -(u_x_d1 **2 + u_y_d1 ** 2) - (u_x_d2 **2 + u_y_d2 ** 2) + \
-            mineral_m1_cond + mineral_m2_cond \
-            - 1. * action['harvest_d1'] - 1. * action['harvest_d2']
+
+    reward = -(u_x_d1 ** 2 + u_y_d1 ** 2) - (u_x_d2 ** 2 + u_y_d2 ** 2) + \
+        mineral_m1_cond + mineral_m2_cond \
+        - 1. * action['harvest_d1'] - 1. * action['harvest_d2']
     return reward, next_state, next_interm
 
+
 def main():
-    myEnv = RDDLEnv.RDDLEnv(domain=FOLDER + 'domain.rddl', instance=FOLDER + 'instance0.rddl', is_grounded=False)
-    agent = RandomAgent(action_space=myEnv.action_space, num_actions=myEnv.NumConcurrentActions)
-    
+    myEnv = RDDLEnv.RDDLEnv(domain=FOLDER + 'domain.rddl',
+                            instance=FOLDER + 'instance0.rddl')
+    agent = RandomAgent(action_space=myEnv.action_space,
+                        num_actions=myEnv.numConcurrentActions)
+
     from pprint import pprint
     pprint(vars(myEnv.model))
     # print(myEnv.model.reward)
-    
+
     print(myEnv.model.cpfs['mineral-harvested_m1\''])
     print(myEnv.model.cpfs['mineral-harvested_m2\''])
     total_reward = 0
@@ -111,13 +118,15 @@ def main():
         print('next_state = {}'.format(next_state))
         print('interm = {}'.format(myEnv.model.interm))
         state = next_state
-        
-        test_reward, next_test_state, test_interm = mars_rover_update(test_state, action)
+
+        test_reward, next_test_state, test_interm = mars_rover_update(
+            test_state, action)
         assert test_reward == reward
         assert next_test_state == next_state
-        assert test_interm == myEnv.model.interm, str(test_interm) + '\n' + str(myEnv.model.interm)
+        assert test_interm == myEnv.model.interm, str(
+            test_interm) + '\n' + str(myEnv.model.interm)
         test_state = next_test_state
-        
+
     print("episode ended with reward {}".format(total_reward))
 
 
diff --git a/pyRDDLGym/tests/unit_test_power_gen.py b/pyRDDLGym/tests/unit_test_power_gen.py
index 950fb381..914f89a6 100644
--- a/pyRDDLGym/tests/unit_test_power_gen.py
+++ b/pyRDDLGym/tests/unit_test_power_gen.py
@@ -1,8 +1,8 @@
 from pyRDDLGym.Core.Env import RDDLEnv as RDDLEnv
-from pyRDDLGym.Policies.Agents import RandomAgent
+from pyRDDLGym.Core.Policies.Agents import RandomAgent
 
 # NOTE: set TEMP_VARIANCE = 0 in RDDL to reproduce result
-FOLDER = 'Competition/Power_gen/'
+FOLDER = 'pyRDDLGym/Examples/PowerGen/Discrete/'
 
 
 def power_gen_update(state, action):
@@ -13,37 +13,39 @@ def power_gen_update(state, action):
     next_prevProd_p2 = prod_p2
     next_prevProd_p3 = prod_p3
     next_temperature = state['temperature']
-    
+
     demand = 2 + 0.01 * (state['temperature'] - 11.7) ** 2
     fulfilledDemand = min(demand, prod_p1 + prod_p2 + prod_p3)
-    
+
     reward = -(prod_p1 * 5 + prod_p2 * 5 + prod_p3 * 5) + (fulfilledDemand * 8) \
-             -(1000. if (demand > fulfilledDemand) else 0.) \
-             +(abs(prod_p1 - state['prevProd_p1']) * 1 + abs(prod_p2 - state['prevProd_p2']) * 1 \
-               +abs(prod_p3 - state['prevProd_p3']) * 1);
-    
+             - (1000. if (demand > fulfilledDemand) else 0.) \
+        + (abs(prod_p1 - state['prevProd_p1']) * 1 + abs(prod_p2 - state['prevProd_p2']) * 1
+           + abs(prod_p3 - state['prevProd_p3']) * 1)
+
     next_interm = {'prod_p1': prod_p1,
-                    'prod_p2': prod_p2,
-                    'prod_p3': prod_p3,
-                    'demand': demand, 
-                    'fulfilledDemand': fulfilledDemand}
-    
+                   'prod_p2': prod_p2,
+                   'prod_p3': prod_p3,
+                   'demand': demand,
+                   'fulfilledDemand': fulfilledDemand}
+
     next_state = {'prevProd_p1': next_prevProd_p1,
                   'prevProd_p2': next_prevProd_p2,
                   'prevProd_p3': next_prevProd_p3,
                   'temperature': next_temperature}
-    
+
     return reward, next_state, next_interm
 
 
 def main():
-    myEnv = RDDLEnv.RDDLEnv(domain=FOLDER + 'domain.rddl', instance=FOLDER + 'instance0.rddl', is_grounded=False)
-    agent = RandomAgent(action_space=myEnv.action_space, num_actions=myEnv.NumConcurrentActions)
+    myEnv = RDDLEnv.RDDLEnv(domain=FOLDER + 'domain.rddl',
+                            instance=FOLDER + 'instance0.rddl')
+    agent = RandomAgent(action_space=myEnv.action_space,
+                        num_actions=myEnv.numConcurrentActions)
     myEnv.model.nonfluents['TEMP_VARIANCE'] = 0.
-    
+
     from pprint import pprint
     pprint(vars(myEnv.model))
-    
+
     total_reward = 0
     state = myEnv.reset()
     test_state = state
@@ -60,13 +62,15 @@ def main():
         print('next_state = {}'.format(next_state))
         print('interm = {}'.format(myEnv.model.interm))
         state = next_state
-        
-        test_reward, next_test_state, test_interm = power_gen_update(test_state, action)
+
+        test_reward, next_test_state, test_interm = power_gen_update(
+            test_state, action)
         assert test_reward == reward
         assert next_test_state == next_state
-        assert test_interm == myEnv.model.interm, str(test_interm) + '\n' + str(myEnv.model.interm)
+        assert test_interm == myEnv.model.interm, str(
+            test_interm) + '\n' + str(myEnv.model.interm)
         test_state = next_test_state
-        
+
     print("episode ended with reward {}".format(total_reward))