code: """Painting domain, which allows for two different grasps on an object (side or
  top).
  
  Side grasping allows for placing into the shelf, and top grasping allows
  for placing into the box. The box has a lid which may need to be opened;
  this lid is NOT modeled by any of the given predicates.
  """
  
  from typing import Any, ClassVar, Dict, List, Optional, Sequence, Set, Tuple, \
      Union
  
  import matplotlib
  import matplotlib.pyplot as plt
  import numpy as np
  from gym.spaces import Box
  from matplotlib import patches
  
  from predicators.src import utils
  from predicators.src.envs import BaseEnv
  from predicators.src.settings import CFG
  from predicators.src.structs import Action, Array, GroundAtom, Object, \
      ParameterizedOption, Predicate, State, Task, Type
  
  
  class PaintingEnv(BaseEnv):
      """Painting domain."""
      # Parameters that aren't important enough to need to clog up settings.py
      table_lb: ClassVar[float] = -10.1
      table_ub: ClassVar[float] = -1.0
      table_height: ClassVar[float] = 0.2
      table_x: ClassVar[float] = 1.65
      shelf_l: ClassVar[float] = 2.0  # shelf length
      shelf_lb: ClassVar[float] = 1.
      shelf_ub: ClassVar[float] = shelf_lb + shelf_l - 0.05
      shelf_x: ClassVar[float] = 1.65
      shelf_y: ClassVar[float] = (shelf_lb + shelf_ub) / 2.0
      box_s: ClassVar[float] = 0.8  # side length
      box_y: ClassVar[float] = 0.5  # y coordinate
      box_lb: ClassVar[float] = box_y - box_s / 10
      box_ub: ClassVar[float] = box_y + box_s / 10
      box_x: ClassVar[float] = 1.65
      env_lb: ClassVar[float] = min(table_lb, shelf_lb, box_lb)
      env_ub: ClassVar[float] = max(table_ub, shelf_ub, box_ub)
      obj_height: ClassVar[float] = 0.13
      obj_radius: ClassVar[float] = 0.03
      obj_x: ClassVar[float] = 1.65
      obj_z: ClassVar[float] = table_height + obj_height / 2
      pick_tol: ClassVar[float] = 1e-2
      color_tol: ClassVar[float] = 1e-2
      wetness_tol: ClassVar[float] = 0.5
      dirtiness_tol: ClassVar[float] = 0.5
      open_fingers: ClassVar[float] = 0.8
      top_grasp_thresh: ClassVar[float] = 0.5 + 1e-2
      side_grasp_thresh: ClassVar[float] = 0.5 - 1e-2
      robot_x: ClassVar[float] = table_x - 0.5
      nextto_thresh: ClassVar[float] = 1.0
      on_table_height_tol: ClassVar[float] = 5e-02
  
      def __init__(self) -> None:
          super().__init__()
          # Types
          self._obj_type = Type("obj", [
              "pose_x", "pose_y", "pose_z", "dirtiness", "wetness", "color",
              "grasp", "held"
          ])
          self._box_type = Type("box", ["pose_x", "pose_y", "color"])
          self._lid_type = Type("lid", ["is_open"])
          self._shelf_type = Type("shelf", ["pose_x", "pose_y", "color"])
          self._robot_type = Type("robot", ["pose_x", "pose_y", "fingers"])
          # Predicates
          self._InBox = Predicate("InBox", [self._obj_type, self._box_type],
                                  self._InBox_holds)
          self._InShelf = Predicate("InShelf",
                                    [self._obj_type, self._shelf_type],
                                    self._InShelf_holds)
          self._IsBoxColor = Predicate("IsBoxColor",
                                       [self._obj_type, self._box_type],
                                       self._IsBoxColor_holds)
          self._IsShelfColor = Predicate("IsShelfColor",
                                         [self._obj_type, self._shelf_type],
                                         self._IsShelfColor_holds)
          self._GripperOpen = Predicate("GripperOpen", [self._robot_type],
                                        self._GripperOpen_holds)
          self._OnTable = Predicate("OnTable", [self._obj_type],
                                    self._OnTable_holds)
          self._NotOnTable = Predicate("NotOnTable", [self._obj_type],
                                       self._NotOnTable_holds)
          self._HoldingTop = Predicate("HoldingTop", [self._obj_type],
                                       self._HoldingTop_holds)
          self._HoldingSide = Predicate("HoldingSide", [self._obj_type],
                                        self._HoldingSide_holds)
          self._Holding = Predicate("Holding", [self._obj_type],
                                    self._Holding_holds)
          self._IsWet = Predicate("IsWet", [self._obj_type], self._IsWet_holds)
          self._IsDry = Predicate("IsDry", [self._obj_type], self._IsDry_holds)
          self._IsDirty = Predicate("IsDirty", [self._obj_type],
                                    self._IsDirty_holds)
          self._IsClean = Predicate("IsClean", [self._obj_type],
                                    self._IsClean_holds)
          # Options
          self._Pick = utils.SingletonParameterizedOption(
              # variables: [robot, object to pick]
              # params: [grasp]
              "Pick",
              self._Pick_policy,
              types=[self._robot_type, self._obj_type],
              params_space=Box(np.array([-0.01], dtype=np.float32),
                               np.array([1.01], dtype=np.float32)))
          self._Wash = utils.SingletonParameterizedOption(
              # variables: [robot]
              # params: []
              "Wash",
              self._Wash_policy,
              types=[self._robot_type])
          self._Dry = utils.SingletonParameterizedOption(
              # variables: [robot]
              # params: []
              "Dry",
              self._Dry_policy,
              types=[self._robot_type])
          self._Paint = utils.SingletonParameterizedOption(
              # variables: [robot]
              # params: [new color]
              "Paint",
              self._Paint_policy,
              types=[self._robot_type],
              params_space=Box(-0.01, 1.01, (1, )))
          self._Place = utils.SingletonParameterizedOption(
              # variables: [robot]
              # params: [absolute x, absolute y, absolute z]
              "Place",
              self._Place_policy,
              types=[self._robot_type],
              params_space=Box(
                  np.array([self.obj_x - 1e-2, self.env_lb, self.obj_z - 1e-2],
                           dtype=np.float32),
                  np.array([self.obj_x + 1e-2, self.env_ub, self.obj_z + 1e-2],
                           dtype=np.float32)))
          self._OpenLid = utils.SingletonParameterizedOption(
              # variables: [robot, lid]
              # params: []
              "OpenLid",
              self._OpenLid_policy,
              types=[self._robot_type, self._lid_type])
          # Static objects (always exist no matter the settings).
          self._box = Object("receptacle_box", self._box_type)
          self._lid = Object("box_lid", self._lid_type)
          self._shelf = Object("receptacle_shelf", self._shelf_type)
          self._robot = Object("robby", self._robot_type)
  
      @classmethod
      def get_name(cls) -> str:
          return "painting"
  
      def simulate(self, state: State, action: Action) -> State:
          assert self.action_space.contains(action.arr)
          arr = action.arr
          # Infer which transition function to follow
          wash_affinity = 0 if arr[5] > 0.5 else abs(arr[5] - 0.5)
          dry_affinity = 0 if arr[6] > 0.5 else abs(arr[6] - 0.5)
          paint_affinity = min(abs(arr[7] - state.get(self._box, "color")),
                               abs(arr[7] - state.get(self._shelf, "color")))
          affinities = [
              (abs(1 - arr[4]), self._transition_pick_or_openlid),
              (wash_affinity, self._transition_wash),
              (dry_affinity, self._transition_dry),
              (paint_affinity, self._transition_paint),
              (abs(-1 - arr[4]), self._transition_place),
          ]
          _, transition_fn = min(affinities, key=lambda item: item[0])
          return transition_fn(state, action)
  
      def _transition_pick_or_openlid(self, state: State,
                                      action: Action) -> State:
          x, y, z, grasp = action.arr[:4]
          next_state = state.copy()
          # Open lid
          if self.box_lb < y < self.box_ub:
              next_state.set(self._lid, "is_open", 1.0)
              return next_state
          held_obj = self._get_held_object(state)
          # Cannot pick if already holding something
          if held_obj is not None:
              return next_state
          # Cannot pick if object pose not on table
          if not self.table_lb < y < self.table_ub:
              return next_state
          # Cannot pick if grasp is invalid
          if self.side_grasp_thresh < grasp < self.top_grasp_thresh:
              return next_state
          # Check if some object is close enough to (x, y, z)
          target_obj = self._get_object_at_xyz(state, x, y, z)
          if target_obj is None:
              return next_state
          # Execute pick
          next_state.set(self._robot, "fingers", 0.0)
          next_state.set(target_obj, "grasp", grasp)
          next_state.set(target_obj, "held", 1.0)
          return next_state
  
      def _transition_wash(self, state: State, action: Action) -> State:
          target_wetness = action.arr[5]
          next_state = state.copy()
          held_obj = self._get_held_object(state)
          # Can only wash if holding obj
          if held_obj is None:
              return next_state
          # Execute wash
          cur_dirtiness = state.get(held_obj, "dirtiness")
          next_dirtiness = max(cur_dirtiness - target_wetness, 0.0)
          next_state.set(held_obj, "wetness", target_wetness)
          next_state.set(held_obj, "dirtiness", next_dirtiness)
          return next_state
  
      def _transition_dry(self, state: State, action: Action) -> State:
          target_wetness = max(1.0 - action.arr[6], 0.0)
          next_state = state.copy()
          held_obj = self._get_held_object(state)
          # Can only dry if holding obj
          if held_obj is None:
              return next_state
          # Execute dry
          next_state.set(held_obj, "wetness", target_wetness)
          return next_state
  
      def _transition_paint(self, state: State, action: Action) -> State:
          color = action.arr[7]
          next_state = state.copy()
          # Can only paint if holding obj
          held_obj = self._get_held_object(state)
          if held_obj is None:
              return next_state
          # Can only paint if dry and clean
          if state.get(held_obj, "dirtiness") > self.dirtiness_tol or \
             state.get(held_obj, "wetness") > self.wetness_tol:
              return next_state
          # Execute paint
          next_state.set(held_obj, "color", color)
          return next_state
  
      def _transition_place(self, state: State, action: Action) -> State:
          # Action args are target pose for held obj
          x, y, z = action.arr[:3]
          next_state = state.copy()
          # Can only place if holding obj
          held_obj = self._get_held_object(state)
          if held_obj is None:
              return next_state
          # Detect table vs shelf vs box place
          if self.table_lb < y < self.table_ub:
              receptacle = "table"
          elif self.shelf_lb < y < self.shelf_ub:
              receptacle = "shelf"
          elif self.box_lb < y < self.box_ub:
              receptacle = "box"
          else:
              # Cannot place outside of table, shelf, or box
              return next_state
          if receptacle == "box" and state.get(self._lid, "is_open") < 0.5:
              # Cannot place in box if lid is not open
              raise utils.EnvironmentFailure("Box lid is closed.",
                                             {"offending_objects": {self._lid}})
          # Detect top grasp vs side grasp
          grasp = state.get(held_obj, "grasp")
          if grasp > self.top_grasp_thresh:
              top_or_side = "top"
          elif grasp < self.side_grasp_thresh:
              top_or_side = "side"
          # Can only place in shelf if side grasping, box if top grasping. If the
          # receptacle is table, we don't care what kind of grasp it is.
          if receptacle == "shelf" and top_or_side != "side":
              return next_state
          if receptacle == "box" and top_or_side != "top":
              return next_state
          # Detect collisions
          collider = self._get_object_at_xyz(state, x, y, z)
          if receptacle == "table" and \
             collider is not None and \
             collider != held_obj:
              return next_state
          # Execute place
          next_state.set(self._robot, "fingers", 1.0)
          next_state.set(held_obj, "pose_x", x)
          next_state.set(held_obj, "pose_y", y)
          if self._update_z_poses:
              if receptacle == "table" and np.allclose(
                      z,
                      self.table_height + self.obj_height / 2,
                      rtol=self.on_table_height_tol):
                  # If placing on table, snap the object to the correct z
                  # position as long as the place location is close enough
                  # (measured by rtol) to the correct table height. This
                  # is necessary for learned samplers to have any hope of
                  # placing objects on the table.
                  next_state.set(held_obj, "pose_z",
                                 self.table_height + self.obj_height / 2)
              else:
                  next_state.set(held_obj, "pose_z", z)
          next_state.set(held_obj, "grasp", 0.5)
          next_state.set(held_obj, "held", 0.0)
          return next_state
  
      @property
      def _num_objects_train(self) -> List[int]:
          return CFG.painting_num_objs_train
  
      @property
      def _num_objects_test(self) -> List[int]:
          return CFG.painting_num_objs_test
  
      def _generate_train_tasks(self) -> List[Task]:
          return self._get_tasks(num_tasks=CFG.num_train_tasks,
                                 num_objs_lst=self._num_objects_train,
                                 rng=self._train_rng)
  
      def _generate_test_tasks(self) -> List[Task]:
          return self._get_tasks(num_tasks=CFG.num_test_tasks,
                                 num_objs_lst=self._num_objects_test,
                                 rng=self._test_rng)
  
      @property
      def predicates(self) -> Set[Predicate]:
          return {
              self._InBox, self._InShelf, self._IsBoxColor, self._IsShelfColor,
              self._GripperOpen, self._OnTable, self._NotOnTable,
              self._HoldingTop, self._HoldingSide, self._Holding, self._IsWet,
              self._IsDry, self._IsDirty, self._IsClean
          }
  
      @property
      def goal_predicates(self) -> Set[Predicate]:
          return {
              self._InBox, self._InShelf, self._IsBoxColor, self._IsShelfColor
          }
  
      @property
      def types(self) -> Set[Type]:
          return {
              self._obj_type, self._box_type, self._lid_type, self._shelf_type,
              self._robot_type
          }
  
      @property
      def options(self) -> Set[ParameterizedOption]:
          return {
              self._Pick, self._Wash, self._Dry, self._Paint, self._Place,
              self._OpenLid
          }
  
      @property
      def action_space(self) -> Box:
          # Actions are 8-dimensional vectors:
          # [x, y, z, grasp, pickplace, water level, heat level, color]
          # Note that pickplace is 1 for pick, -1 for place, and 0 otherwise,
          # while grasp, water level, heat level, and color are in [0, 1].
          # We set the lower bound for z to 0.0, rather than self.obj_z - 1e-2,
          # because in RepeatedNextToPainting, we use this dimension to check
          # affinity of the move action
          lowers = np.array(
              [self.obj_x - 1e-2, self.env_lb, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0],
              dtype=np.float32)
          uppers = np.array([
              self.obj_x + 1e-2, self.env_ub, self.obj_z + 1e-2, 1.0, 1.0, 1.0,
              1.0, 1.0
          ],
                            dtype=np.float32)
          return Box(lowers, uppers)
  
      def render_state_plt(
              self,
              state: State,
              task: Task,
              action: Optional[Action] = None,
              caption: Optional[str] = None) -> matplotlib.figure.Figure:
          fig, ax = plt.subplots(1, 1)
          objs = [o for o in state if o.is_instance(self._obj_type)]
          denom = (self.env_ub - self.env_lb)
          # The factor of "2" here should actually be 0.5, but this
          # makes the objects too small, so we'll let it be bigger.
          # Don't be alarmed if objects seem to be intersecting in
          # the resulting videos.
          r = 2 * self.obj_radius / denom
          h = 2 * self.obj_height / denom
          z = (self.obj_z - self.env_lb) / denom
          # Draw box
          box_color = state.get(self._box, "color")
          box_lower = (self.box_lb - self.obj_radius - self.env_lb) / denom
          box_upper = (self.box_ub + self.obj_radius - self.env_lb) / denom
          rect = plt.Rectangle((box_lower, z - h),
                               box_upper - box_lower,
                               2 * h,
                               facecolor=[box_color, 0, 0],
                               alpha=0.25)
          ax.add_patch(rect)
          # Draw box lid
          if state.get(self._lid, "is_open") < 0.5:
              plt.plot([box_lower, box_upper], [z + h, z + h],
                       color=[box_color, 0, 0])
          # Draw shelf
          shelf_color = state.get(self._shelf, "color")
          shelf_lower = (self.shelf_lb - self.obj_radius - self.env_lb) / denom
          shelf_upper = (self.shelf_ub + self.obj_radius - self.env_lb) / denom
          rect = plt.Rectangle((shelf_lower, z - h),
                               shelf_upper - shelf_lower,
                               2 * h,
                               facecolor=[shelf_color, 0, 0],
                               alpha=0.25)
          ax.add_patch(rect)
          # Draw objects
          held_obj = self._get_held_object(state)
          for obj in sorted(objs):
              x = state.get(obj, "pose_x")
              y = state.get(obj, "pose_y")
              z = state.get(obj, "pose_z")
              facecolor: Union[None, str, List[Any]] = None
              if state.get(obj, "wetness") > self.wetness_tol and \
                 state.get(obj, "dirtiness") < self.dirtiness_tol:
                  # wet and clean
                  facecolor = "blue"
              elif state.get(obj, "wetness") < self.wetness_tol and \
                   state.get(obj, "dirtiness") > self.dirtiness_tol:
                  # dry and dirty
                  facecolor = "green"
              elif state.get(obj, "wetness") < self.wetness_tol and \
                   state.get(obj, "dirtiness") < self.dirtiness_tol:
                  # dry and clean
                  facecolor = "cyan"
              obj_color = state.get(obj, "color")
              if obj_color > 0:
                  facecolor = [obj_color, 0, 0]
              if held_obj == obj:
                  assert state.get(self._robot, "fingers") < self.open_fingers
                  grasp = state.get(held_obj, "grasp")
                  assert grasp < self.side_grasp_thresh or \
                      grasp > self.top_grasp_thresh
                  edgecolor = ("yellow"
                               if grasp < self.side_grasp_thresh else "orange")
              else:
                  edgecolor = "gray"
              # Normalize poses to [0, 1]
              x = (x - self.env_lb) / denom
              y = (y - self.env_lb) / denom
              z = (z - self.env_lb) / denom
              # Plot as rectangle
              rect = patches.Rectangle((y - r, z - h),
                                       2 * r,
                                       2 * h,
                                       zorder=-x,
                                       linewidth=1,
                                       edgecolor=edgecolor,
                                       facecolor=facecolor)
              ax.add_patch(rect)
          ax.set_xlim(-0.1, 1.1)
          ax.set_ylim(0.6, 1.0)
          title = ("blue = wet+clean, green = dry+dirty, cyan = dry+clean;\n"
                   "yellow border = side grasp, orange border = top grasp")
          if caption is not None:
              title += f";\n{caption}"
          plt.suptitle(title, fontsize=12, wrap=True)
          plt.tight_layout()
          return fig
  
      @property
      def _max_objs_in_goal(self) -> int:
          return CFG.painting_max_objs_in_goal
  
      @property
      def _update_z_poses(self) -> bool:
          return False
  
      def _get_tasks(self, num_tasks: int, num_objs_lst: List[int],
                     rng: np.random.Generator) -> List[Task]:
          tasks = []
          for i in range(num_tasks):
              num_objs = num_objs_lst[i % len(num_objs_lst)]
              data = {}
              # Initialize robot pos with open fingers
              robot_init_y = rng.uniform(self.table_lb, self.table_ub)
              data[self._robot] = np.array([self.robot_x, robot_init_y, 1.0],
                                           dtype=np.float32)
              # Sample distinct colors for shelf and box
              color1 = rng.uniform(0.2, 0.4)
              color2 = rng.uniform(0.6, 1.0)
              if rng.choice(2):
                  box_color, shelf_color = color1, color2
              else:
                  shelf_color, box_color = color1, color2
              # Create box, lid, and shelf objects
              lid_is_open = int(rng.uniform() < CFG.painting_lid_open_prob)
              data[self._box] = np.array([self.box_x, self.box_y, box_color],
                                         dtype=np.float32)
              data[self._lid] = np.array([lid_is_open], dtype=np.float32)
              data[self._shelf] = np.array(
                  [self.shelf_x, self.shelf_y, shelf_color], dtype=np.float32)
              # Create moveable objects and goal
              objs = []
              obj_poses: List[Tuple[float, float, float]] = []
              goal = set()
              assert CFG.painting_goal_receptacles in ("box_and_shelf", "box",
                                                       "shelf")
              if CFG.painting_goal_receptacles == "shelf":
                  # No box; all max_objs_in_goal objects must go in the shelf
                  num_objs_in_shelf = self._max_objs_in_goal
              else:
                  # The last object is destined for the box, so the remaining
                  # (max_objs_in_goal - 1) objects must go in the shelf
                  num_objs_in_shelf = self._max_objs_in_goal - 1
              for j in range(num_objs):
                  obj = Object(f"obj{j}", self._obj_type)
                  objs.append(obj)
                  pose = self._sample_initial_object_pose(obj_poses, rng)
                  obj_poses.append(pose)
                  # Start out wet and clean, dry and dirty, or dry and clean
                  choice = rng.choice(3)
                  if choice == 0:
                      wetness = 0.0
                      dirtiness = rng.uniform(0.5, 1.)
                  elif choice == 1:
                      wetness = rng.uniform(0.5, 1.)
                      dirtiness = 0.0
                  else:
                      wetness = 0.0
                      dirtiness = 0.0
                  color = 0.0
                  grasp = 0.5
                  held = 0.0
                  data[obj] = np.array([
                      pose[0], pose[1], pose[2], dirtiness, wetness, color,
                      grasp, held
                  ],
                                       dtype=np.float32)
                  if CFG.painting_goal_receptacles in (
                          "box_and_shelf", "box") and j == num_objs - 1:
                      # This object must go in the box
                      # NOTE: the box can only fit one object
                      goal.add(GroundAtom(self._InBox, [obj, self._box]))
                      goal.add(GroundAtom(self._IsBoxColor, [obj, self._box]))
                  elif CFG.painting_goal_receptacles in (
                          "box_and_shelf", "shelf") and j < num_objs_in_shelf:
                      # This object must go in the shelf
                      # NOTE: any number of objects can fit in the shelf
                      goal.add(GroundAtom(self._InShelf, [obj, self._shelf]))
                      goal.add(GroundAtom(self._IsShelfColor,
                                          [obj, self._shelf]))
              assert len(goal) <= 2 * self._max_objs_in_goal
              state = State(data)
              # Sometimes start out holding an object, possibly with the wrong
              # grip, so that we'll have to put it on the table and regrasp
              if rng.uniform() < CFG.painting_initial_holding_prob:
                  grasp = rng.choice([0.0, 1.0])
                  target_obj = objs[rng.choice(len(objs))]
                  state.set(self._robot, "fingers", 0.0)
                  state.set(target_obj, "grasp", grasp)
                  state.set(target_obj, "held", 1.0)
                  state.set(target_obj, "pose_y",
                            state.get(self._robot, "pose_y"))
                  if self._update_z_poses:
                      state.set(target_obj, "pose_z",
                                state.get(target_obj, "pose_z") + 1.0)
              tasks.append(Task(state, goal))
          return tasks
  
      def _sample_initial_object_pose(
              self, existing_poses: List[Tuple[float, float, float]],
              rng: np.random.Generator) -> Tuple[float, float, float]:
          existing_ys = [p[1] for p in existing_poses]
          while True:
              this_y = rng.uniform(self.table_lb, self.table_ub)
              if all(
                      abs(this_y - other_y) > 3.5 * self.obj_radius
                      for other_y in existing_ys):
                  return (self.obj_x, this_y,
                          self.table_height + self.obj_height / 2)
  
      def _Pick_policy(self, state: State, memory: Dict,
                       objects: Sequence[Object], params: Array) -> Action:
          del memory  # unused
          _, obj = objects
          obj_x = state.get(obj, "pose_x")
          obj_y = state.get(obj, "pose_y")
          obj_z = state.get(obj, "pose_z")
          grasp, = params
          arr = np.array([obj_x, obj_y, obj_z, grasp, 1.0, 0.0, 0.0, 0.0],
                         dtype=np.float32)
          # The grasp could cause the action to go out of bounds, so we clip
          # it back into the bounds for safety.
          arr = np.clip(arr, self.action_space.low, self.action_space.high)
          return Action(arr)
  
      def _Wash_policy(self, state: State, memory: Dict,
                       objects: Sequence[Object], params: Array) -> Action:
          del state, memory, objects, params  # unused
          arr = np.array(
              [self.obj_x, self.table_lb, self.obj_z, 0.0, 0.0, 1.0, 0.0, 0.0],
              dtype=np.float32)
          return Action(arr)
  
      def _Dry_policy(self, state: State, memory: Dict,
                      objects: Sequence[Object], params: Array) -> Action:
          del state, memory, objects, params  # unused
          arr = np.array(
              [self.obj_x, self.table_lb, self.obj_z, 0.0, 0.0, 0.0, 1.0, 0.0],
              dtype=np.float32)
          return Action(arr)
  
      def _Paint_policy(self, state: State, memory: Dict,
                        objects: Sequence[Object], params: Array) -> Action:
          del state, memory, objects  # unused
          new_color, = params
          new_color = min(max(new_color, 0.0), 1.0)
          arr = np.array([
              self.obj_x, self.table_lb, self.obj_z, 0.0, 0.0, 0.0, 0.0,
              new_color
          ],
                         dtype=np.float32)
          return Action(arr)
  
      @staticmethod
      def _Place_policy(state: State, memory: Dict, objects: Sequence[Object],
                        params: Array) -> Action:
          del state, memory, objects  # unused
          x, y, z = params
          arr = np.array([x, y, z, 0.5, -1.0, 0.0, 0.0, 0.0], dtype=np.float32)
          return Action(arr)
  
      def _OpenLid_policy(self, state: State, memory: Dict,
                          objects: Sequence[Object], params: Array) -> Action:
          del state, memory, objects, params  # unused
          arr = np.array([
              self.obj_x, (self.box_lb + self.box_ub) / 2, self.obj_z, 0.0, 1.0,
              0.0, 0.0, 0.0
          ],
                         dtype=np.float32)
          return Action(arr)
  
      def _InBox_holds(self, state: State, objects: Sequence[Object]) -> bool:
          obj, _ = objects
          # If the object is held, not yet in box
          if self._obj_is_held(state, obj):
              return False
          # Check pose of object
          obj_y = state.get(obj, "pose_y")
          return self.box_lb < obj_y < self.box_ub
  
      def _InShelf_holds(self, state: State, objects: Sequence[Object]) -> bool:
          obj, _ = objects
          # If the object is held, not yet in shelf
          if self._obj_is_held(state, obj):
              return False
          # Check pose of object
          obj_y = state.get(obj, "pose_y")
          return self.shelf_lb < obj_y < self.shelf_ub
  
      def _IsBoxColor_holds(self, state: State,
                            objects: Sequence[Object]) -> bool:
          obj, box = objects
          return abs(state.get(obj, "color") -
                     state.get(box, "color")) < self.color_tol
  
      def _IsShelfColor_holds(self, state: State,
                              objects: Sequence[Object]) -> bool:
          obj, shelf = objects
          return abs(state.get(obj, "color") -
                     state.get(shelf, "color")) < self.color_tol
  
      def _GripperOpen_holds(self, state: State,
                             objects: Sequence[Object]) -> bool:
          robot, = objects
          fingers = state.get(robot, "fingers")
          return fingers >= self.open_fingers
  
      def _OnTable_holds(self, state: State, objects: Sequence[Object]) -> bool:
          obj, = objects
          obj_y = state.get(obj, "pose_y")
          if not self.table_lb < obj_y < self.table_ub:
              return False
          # Note that obj_z is not updated in this class, but it may be updated
          # by subclasses by overriding self._update_z_poses.
          obj_z = state.get(obj, "pose_z")
          if not np.allclose(obj_z, self.table_height + self.obj_height / 2):
              assert self._update_z_poses
              return False
          return True
  
      def _NotOnTable_holds(self, state: State,
                            objects: Sequence[Object]) -> bool:
          return not self._OnTable_holds(state, objects)
  
      def _HoldingTop_holds(self, state: State,
                            objects: Sequence[Object]) -> bool:
          obj, = objects
          grasp = state.get(obj, "grasp")
          return grasp > self.top_grasp_thresh
  
      def _HoldingSide_holds(self, state: State,
                             objects: Sequence[Object]) -> bool:
          obj, = objects
          grasp = state.get(obj, "grasp")
          return grasp < self.side_grasp_thresh
  
      def _Holding_holds(self, state: State, objects: Sequence[Object]) -> bool:
          obj, = objects
          return self._obj_is_held(state, obj)
  
      def _IsWet_holds(self, state: State, objects: Sequence[Object]) -> bool:
          obj, = objects
          return state.get(obj, "wetness") > self.wetness_tol
  
      def _IsDry_holds(self, state: State, objects: Sequence[Object]) -> bool:
          obj, = objects
          return not self._IsWet_holds(state, [obj])
  
      def _IsDirty_holds(self, state: State, objects: Sequence[Object]) -> bool:
          obj, = objects
          return state.get(obj, "dirtiness") > self.dirtiness_tol
  
      def _IsClean_holds(self, state: State, objects: Sequence[Object]) -> bool:
          obj, = objects
          return not self._IsDirty_holds(state, [obj])
  
      def _get_held_object(self, state: State) -> Optional[Object]:
          for obj in state:
              if obj.type != self._obj_type:
                  continue
              if self._obj_is_held(state, obj):
                  return obj
          return None
  
      def _obj_is_held(self, state: State, obj: Object) -> bool:
          # These two pieces of information are redundant. We include
          # the "held" feature only because it allows the Holding
          # predicate to be expressed with a single inequality.
          # Either feature can be used to implement this method.
          grasp = state.get(obj, "grasp")
          held_feat = state.get(obj, "held")
          is_held = (grasp > self.top_grasp_thresh
                     or grasp < self.side_grasp_thresh)
          assert is_held == (held_feat > 0.5)  # ensure redundancy
          return is_held
  
      def _get_object_at_xyz(self, state: State, x: float, y: float,
                             z: float) -> Optional[Object]:
          target_obj = None
          for obj in state:
              if obj.type != self._obj_type:
                  continue
              if np.allclose([x, y, z], [
                      state.get(obj, "pose_x"),
                      state.get(obj, "pose_y"),
                      state.get(obj, "pose_z")
              ],
                             atol=self.pick_tol):
                  target_obj = obj
          return target_obj
  
apis: numpy.allclose matplotlib.pyplot.tight_layout matplotlib.pyplot.subplots numpy.clip matplotlib.patches.Rectangle matplotlib.pyplot.suptitle numpy.array matplotlib.pyplot.plot matplotlib.pyplot.Rectangle