Add env-specific process samplers for continuous option params

Replace null_sampler with proper samplers that return fixed param values
for each env's EndogenousProcess (pick, place, push, pour). Place
samplers read target coordinates from state or process objects.

Author: yichao-liang

predicators/ground_truth_models/boil/processes.py

@@ -1,20 +1,70 @@
 """Ground-truth processes for the boil environments."""
 import logging
 from pprint import pformat
-from typing import Dict, Set, cast
+from typing import Dict, Sequence, Set, cast
 
 import numpy as np
 import torch
 
+from predicators.envs.pybullet_boil import PyBulletBoilEnv
 from predicators.ground_truth_models import GroundTruthProcessFactory
 from predicators.settings import CFG
-from predicators.structs import CausalProcess, DelayDistribution, \
-    EndogenousProcess, ExogenousProcess, LiftedAtom, ParameterizedOption, \
-    Predicate, Type, Variable
+from predicators.structs import Array, CausalProcess, DelayDistribution, \
+    EndogenousProcess, ExogenousProcess, GroundAtom, LiftedAtom, Object, \
+    ParameterizedOption, Predicate, State, Type, Variable
 from predicators.utils import ConstantDelay, DiscreteGaussianDelay, \
     null_sampler
 
 
+_BOIL_DROP_Z = 0.49  # table_height (0.4) + jug_handle_height (0.09)
+
+
+def _pick_sampler(state: State, goal: Set[GroundAtom],
+                  rng: np.random.Generator,
+                  objs: Sequence[Object]) -> Array:
+    del state, goal, rng, objs
+    return np.array([0.0], dtype=np.float32)
+
+
+def _push_sampler(state: State, goal: Set[GroundAtom],
+                  rng: np.random.Generator,
+                  objs: Sequence[Object]) -> Array:
+    del state, goal, rng, objs
+    return np.array([0.057, 0.104, 0.0, 0.25], dtype=np.float32)
+
+
+def _place_on_burner_sampler(state: State, goal: Set[GroundAtom],
+                              rng: np.random.Generator,
+                              objs: Sequence[Object]) -> Array:
+    del goal, rng
+    # objs = [robot, jug, burner]
+    burner = objs[2]
+    x = state.get(burner, "x")
+    y = state.get(burner, "y") - PyBulletBoilEnv.jug_handle_offset
+    return np.array([x, y, 0.0, _BOIL_DROP_Z], dtype=np.float32)
+
+
+def _place_under_faucet_sampler(state: State, goal: Set[GroundAtom],
+                                 rng: np.random.Generator,
+                                 objs: Sequence[Object]) -> Array:
+    del goal, rng
+    # objs = [robot, jug, faucet]
+    faucet = objs[2]
+    x = state.get(faucet, "x")
+    y = (state.get(faucet, "y") - PyBulletBoilEnv.jug_handle_offset
+         - PyBulletBoilEnv.faucet_x_len)
+    return np.array([x, y, 0.0, _BOIL_DROP_Z], dtype=np.float32)
+
+
+def _place_outside_sampler(state: State, goal: Set[GroundAtom],
+                            rng: np.random.Generator,
+                            objs: Sequence[Object]) -> Array:
+    del state, goal, rng, objs
+    x = PyBulletBoilEnv.x_mid - 0.15
+    y = PyBulletBoilEnv.y_mid + 0.10
+    return np.array([x, y, 0.0, _BOIL_DROP_Z], dtype=np.float32)
+
+
 class PyBulletBoilGroundTruthProcessFactory(GroundTruthProcessFactory):
     """Ground-truth processes for the boil environment."""
 
@@ -107,7 +157,7 @@ def get_processes(
         pick_jug_from_faucet_process = EndogenousProcess(
             "PickJugFromFaucet", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _pick_sampler)
         processes.add(pick_jug_from_faucet_process)
 
         # PickJugFromBurner
@@ -134,7 +184,7 @@ def get_processes(
         pick_jug_from_burner_process = EndogenousProcess(
             "PickJugFromBurner", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _pick_sampler)
         processes.add(pick_jug_from_burner_process)
 
         # PickJugFromOutsideFaucetAndBurner
@@ -160,7 +210,7 @@ def get_processes(
             "PickJugFromOutsideFaucetAndBurner", parameters,
             condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _pick_sampler)
         processes.add(pick_jug_outside_faucet_burner_process)
 
         # PlaceOnBurner
@@ -187,7 +237,7 @@ def get_processes(
         place_on_burner_process = EndogenousProcess(
             "PlaceOnBurner", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _place_on_burner_sampler)
         processes.add(place_on_burner_process)
 
         # PlaceUnderFaucet
@@ -214,7 +264,8 @@ def get_processes(
         place_under_faucet_process = EndogenousProcess(
             "PlaceUnderFaucet", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars,
+            _place_under_faucet_sampler)
         processes.add(place_under_faucet_process)
 
         # PlaceAtOutsideFaucetAndBurner
@@ -238,7 +289,7 @@ def get_processes(
         place_at_outside_faucet_burner_process = EndogenousProcess(
             "PlaceOutsideFaucetAndBurner", parameters, condition_at_start,
             set(), set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _place_outside_sampler)
         processes.add(place_at_outside_faucet_burner_process)
 
         # SwitchFaucetOn
@@ -262,7 +313,7 @@ def get_processes(
         switch_faucet_on_process = EndogenousProcess(
             "SwitchFaucetOn", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _push_sampler)
         processes.add(switch_faucet_on_process)
 
         # SwitchFaucetOff
@@ -286,7 +337,7 @@ def get_processes(
         switch_faucet_off_process = EndogenousProcess(
             "SwitchFaucetOff", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _push_sampler)
         processes.add(switch_faucet_off_process)
 
         # SwitchBurnerOn
@@ -310,7 +361,7 @@ def get_processes(
         switch_burner_on_process = EndogenousProcess(
             "SwitchBurnerOn", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _push_sampler)
         processes.add(switch_burner_on_process)
 
         # SwitchBurnerOff
@@ -334,7 +385,7 @@ def get_processes(
         switch_burner_off_process = EndogenousProcess(
             "SwitchBurnerOff", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _push_sampler)
         processes.add(switch_burner_off_process)
 
         # Noop

predicators/ground_truth_models/coffee/processes.py

@@ -1,18 +1,54 @@
 """Ground-truth processes for the coffee environments."""
 
-from typing import Dict, Set, cast
+from typing import Dict, Sequence, Set, cast
 
 import numpy as np
 import torch
 
+from predicators.envs.pybullet_coffee import PyBulletCoffeeEnv
 from predicators.ground_truth_models import GroundTruthProcessFactory
 from predicators.settings import CFG
-from predicators.structs import CausalProcess, DelayDistribution, \
-    EndogenousProcess, ExogenousProcess, LiftedAtom, ParameterizedOption, \
-    Predicate, Type, Variable
+from predicators.structs import Array, CausalProcess, DelayDistribution, \
+    EndogenousProcess, ExogenousProcess, GroundAtom, LiftedAtom, Object, \
+    ParameterizedOption, Predicate, State, Type, Variable
 from predicators.utils import ConstantDelay, DiscreteGaussianDelay, \
     null_sampler
 
+_COFFEE_DROP_Z = 0.5  # z_lb (0.4) + jug_handle_height (0.1)
+
+
+def _pick_sampler(state: State, goal: Set[GroundAtom],
+                  rng: np.random.Generator,
+                  objs: Sequence[Object]) -> Array:
+    del state, goal, rng, objs
+    return np.array([0.0], dtype=np.float32)
+
+
+def _push_sampler(state: State, goal: Set[GroundAtom],
+                  rng: np.random.Generator,
+                  objs: Sequence[Object]) -> Array:
+    """Push params for TurnMachineOn (button press)."""
+    del state, goal, rng, objs
+    return np.array([0.0675, 0.0, -np.pi, 0.0], dtype=np.float32)
+
+
+def _place_jug_in_machine_sampler(state: State, goal: Set[GroundAtom],
+                                   rng: np.random.Generator,
+                                   objs: Sequence[Object]) -> Array:
+    del state, goal, rng
+    # objs = [robot, jug, machine]
+    return np.array([PyBulletCoffeeEnv.dispense_area_x,
+                     PyBulletCoffeeEnv.dispense_area_y,
+                     PyBulletCoffeeEnv.robot_init_wrist,
+                     _COFFEE_DROP_Z], dtype=np.float32)
+
+
+def _pour_sampler(state: State, goal: Set[GroundAtom],
+                  rng: np.random.Generator,
+                  objs: Sequence[Object]) -> Array:
+    del state, goal, rng, objs
+    return np.array([np.pi / 4], dtype=np.float32)
+
 
 class PyBulletCoffeeGroundTruthProcessFactory(GroundTruthProcessFactory):
     """Ground-truth processes for the coffee environment."""
@@ -189,7 +225,7 @@ def get_processes(
         pick_jug_from_table_process = EndogenousProcess(
             "PickJugFromTable", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _pick_sampler)
         processes.add(pick_jug_from_table_process)
 
         # PlaceJugInMachine
@@ -215,7 +251,8 @@ def get_processes(
         place_jug_in_machine_process = EndogenousProcess(
             "PlaceJugInMachine", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars,
+            _place_jug_in_machine_sampler)
         processes.add(place_jug_in_machine_process)
 
         # TurnMachineOn
@@ -242,7 +279,7 @@ def get_processes(
         turn_machine_on_process = EndogenousProcess(
             "TurnMachineOn", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _push_sampler)
         processes.add(turn_machine_on_process)
 
         # PickJugFromMachine
@@ -268,7 +305,7 @@ def get_processes(
         pick_jug_from_machine_process = EndogenousProcess(
             "PickJugFromMachine", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _pick_sampler)
         processes.add(pick_jug_from_machine_process)
 
         # Pour from not-above-cup
@@ -293,7 +330,7 @@ def get_processes(
         pourFromNotAboveCup_process = EndogenousProcess(
             "PourFromNotAboveCup", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler)
+            torch.tensor(1.0), option, option_vars, _pour_sampler)
         processes.add(pourFromNotAboveCup_process)
 
         # Pour from above-cup
@@ -321,7 +358,7 @@ def get_processes(
         pourFromNotAboveCup_process = EndogenousProcess(
             "PourFromCup", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler,
+            torch.tensor(1.0), option, option_vars, _pour_sampler,
             ignore_effects)
         processes.add(pourFromNotAboveCup_process)
 

predicators/ground_truth_models/domino/processes.py

@@ -1,17 +1,59 @@
 """Ground-truth processes for the domino environment."""
 
-from typing import Dict, Set
+from typing import Dict, Sequence, Set
 
+import numpy as np
 import torch
 
 from predicators.ground_truth_models import GroundTruthProcessFactory
 from predicators.settings import CFG
-from predicators.structs import CausalProcess, EndogenousProcess, \
-    ExogenousProcess, LiftedAtom, ParameterizedOption, Predicate, Type, \
-    Variable
+from predicators.structs import Array, CausalProcess, EndogenousProcess, \
+    ExogenousProcess, GroundAtom, LiftedAtom, Object, ParameterizedOption, \
+    Predicate, State, Type, Variable
 from predicators.utils import ConstantDelay, DiscreteGaussianDelay, \
     null_sampler
 
+# Fixed parameter values for domino environment.
+_DOMINO_GRASP_Z_OFFSET = 0.0825  # domino_height * 0.55
+_DOMINO_DROP_Z = 0.5695  # table_height + domino_height * 1.13
+_DOMINO_OFFSET_X = 0.045  # domino_depth * 3
+_DOMINO_OFFSET_Z = 0.0825  # domino_height * 0.55
+_DOMINO_OFFSET_ROT = np.pi / 2
+_DOMINO_PUSH_THROUGH_FRAC = 0.25
+
+
+def _pick_sampler(state: State, goal: Set[GroundAtom],
+                  rng: np.random.Generator,
+                  objs: Sequence[Object]) -> Array:
+    """Return fixed grasp_z_offset for domino pick."""
+    del state, goal, rng, objs
+    return np.array([_DOMINO_GRASP_Z_OFFSET], dtype=np.float32)
+
+
+def _push_sampler(state: State, goal: Set[GroundAtom],
+                  rng: np.random.Generator,
+                  objs: Sequence[Object]) -> Array:
+    """Return fixed push params for domino push."""
+    del state, goal, rng, objs
+    return np.array([_DOMINO_OFFSET_X, _DOMINO_OFFSET_Z,
+                     _DOMINO_OFFSET_ROT, _DOMINO_PUSH_THROUGH_FRAC],
+                    dtype=np.float32)
+
+
+def _place_sampler(state: State, goal: Set[GroundAtom],
+                   rng: np.random.Generator,
+                   objs: Sequence[Object]) -> Array:
+    """Return placement params from process objects."""
+    del state, goal, rng
+    # objs = [robot, domino1, domino2, target_pos, rotation]
+    target_pos = objs[3]
+    rotation = objs[4]
+    x = float(target_pos.name.split("_")[1])
+    y = float(target_pos.name.split("_")[2])
+    angle_deg = float(rotation.name.split("_")[-1])
+    yaw = np.radians(angle_deg)
+    return np.array([x, y, yaw, _DOMINO_DROP_Z], dtype=np.float32)
+
 
 class PyBulletDominoGroundTruthProcessFactory(GroundTruthProcessFactory):
     """Ground-truth processes for the domino grid environment."""
@@ -86,7 +128,7 @@ def get_processes(
         push_start_block_process = EndogenousProcess(
             "PushStartBlock", parameters, condition_at_start, set(),
             set(), add_effects, delete_effects, delay_distribution,
-            torch.tensor(1.0), option, option_vars, null_sampler,
+            torch.tensor(1.0), option, option_vars, _push_sampler,
             ignore_effects)
         processes.add(push_start_block_process)
 
@@ -125,7 +167,7 @@ def get_processes(
                                                 delete_effects,
                                                 delay_distribution,
                                                 torch.tensor(1.0), option,
-                                                option_vars, null_sampler,
+                                                option_vars, _pick_sampler,
                                                 ignore_effects)
         processes.add(pick_domino_process)
 
@@ -137,7 +179,7 @@ def get_processes(
         target_pos = Variable("?pos1", position_type)
         rotation = Variable("?rot", rotation_type)
         parameters = [robot, domino1, domino2, target_pos, rotation]
-        option_vars = [robot, domino1, domino2, target_pos, rotation]
+        option_vars = [robot]
         option = Place
         condition_at_start = {
             LiftedAtom(Holding, [robot, domino1]),
@@ -165,7 +207,7 @@ def get_processes(
                                                  delete_effects,
                                                  delay_distribution,
                                                  torch.tensor(1.0), option,
-                                                 option_vars, null_sampler,
+                                                 option_vars, _place_sampler,
                                                  ignore_effects)
         processes.add(place_domino_process)