tabletask 0.2.0__tar.gz → 0.2.1__tar.gz

{tabletask-0.2.0 → tabletask-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: tabletask
-Version: 0.2.0
+Version: 0.2.1
 Summary: Tabletop manipulation tasks as JSON, sampled and graded in MuJoCo
 Project-URL: Homepage, https://table-task.com
 Project-URL: Source, https://github.com/omarrayyann/tabletask

{tabletask-0.2.0 → tabletask-0.2.1}/tabletask/__init__.py

@@ -10,7 +10,7 @@ from .language.schema import Task
 from .sampling import instantiate
 from .vector import Envs
 
-__version__ = "0.2.0"
+__version__ = "0.2.1"
 
 __all__ = [
     "Env",
@@ -38,6 +38,7 @@ def make(
     robot="droid",
     gpus=None,
     renderer=None,
+    reward="success",
     scene=None,
     bindings=None,
     monitor_interval=0.01,
@@ -52,6 +53,7 @@ def make(
         robot: The robot that acts in the scene: "droid", "gripper", a `Robot`, or None.
         gpus: With `num_envs`, CUDA GPU indices the environments render on, in turn.
         renderer: None to rasterise images, or "rtx" to path-trace them.
+        reward: "success" for 1.0 on success, or "subtasks" for 1.0 at each new milestone.
         scene: A scene XML file to use in place of the task's own.
         bindings: Assets for abstract objects, as {name: {"asset": path, "properties": {...}}}.
         monitor_interval: How often, in simulated seconds, task conditions are checked.
@@ -67,6 +69,7 @@ def make(
             robot=robot,
             gpus=gpus,
             renderer=renderer,
+            reward=reward,
             scene=scene,
             bindings=bindings,
             monitor_interval=monitor_interval,
@@ -76,6 +79,7 @@ def make(
         raise ValueError("gpus go with num_envs")
     env = Env(load(source), scene, bindings, monitor_interval=monitor_interval, backend=backend)
     env.renderer = renderer
+    env.reward = reward
     if isinstance(robot, str):
         robot = _robot(robot)
     return env if robot is None else env.add(robot)

{tabletask-0.2.0 → tabletask-0.2.1}/tabletask/env.py

@@ -26,6 +26,7 @@ class Env:
         end_on_fall: Whether an object falling off the table ends the episode.
         image_size: Observation image size as (width, height).
         renderer: None to rasterise images, or "rtx" to path-trace them.
+        reward: "success" or "subtasks", as `step` describes.
     """
 
     def __init__(
@@ -45,6 +46,7 @@ class Env:
         self.end_on_fall = True
         self.image_size = (640, 368)  # 16:9, the shape of DROID's camera frames
         self.renderer = None
+        self.reward = "success"
 
     @property
     def monitor_interval(self) -> float:
@@ -84,6 +86,8 @@ class Env:
         """
         if self.renderer not in (None, "rtx"):
             raise ValueError(f"unknown renderer {self.renderer!r}; use None or 'rtx'")
+        if self.reward not in ("success", "subtasks"):
+            raise ValueError(f"unknown reward {self.reward!r}; use 'success' or 'subtasks'")
         self.close()
         episode = episode if seed is None else seed
         if self.robot is not None and hasattr(self.robot, "start_pose"):
@@ -103,6 +107,7 @@ class Env:
                 self.runtime.use_rtx(*self.image_size)
             self.monitor = Monitor(self.task)
             self.status = self.monitor.reset(self.runtime)
+            self._best = self.status.progress
             return self.observe(), self._info()
         except BaseException:
             self.close()
@@ -120,8 +125,10 @@ class Env:
 
         Returns:
             (observation, reward, terminated, truncated, info). The reward is 1.0 once
-            the task has succeeded. An episode terminates on success or on a broken
-            always condition, and is truncated at the task's `time_limit`, if it has one.
+            the task has succeeded; with `reward` set to "subtasks", it is 1.0 at each step
+            that first reaches a new milestone (see `Monitor`), and never paid twice. An
+            episode terminates on success or on a broken always condition, and is truncated
+            at the task's `time_limit`, if it has one.
         """
         self._require_reset()
         if self.robot is None and action is not None:
@@ -147,7 +154,11 @@ class Env:
         succeeded = self.status.succeeded(self.task)
         terminated = succeeded or bool(self.status.violated)
         truncated = bool(self.status.timed_out)
-        return self.observe(), float(succeeded), terminated, truncated, self._info()
+        reward = float(succeeded)
+        if self.reward == "subtasks":
+            reward = float(max(self.status.progress - self._best, 0))
+            self._best = max(self._best, self.status.progress)
+        return self.observe(), reward, terminated, truncated, self._info()
 
     def _info(self) -> dict:
         return {

{tabletask-0.2.0 → tabletask-0.2.1}/tabletask/language/monitor.py

@@ -8,7 +8,7 @@ from dataclasses import dataclass, field
 from .evaluation import evaluate, observe_history, unpack_expression
 from .modules import require_capabilities
 from .schema import Task
-from .transformations import concrete
+from .transformations import concrete, concrete_condition
 from .vocabulary import CONNECTIVES
 
 # An object that moved more than this between two checks is still moving.
@@ -24,6 +24,7 @@ class Status:
     done: list[str] = field(default_factory=list)
     violated: str = ""
     timed_out: bool = False
+    progress: int = 0
 
     @property
     def failed(self) -> bool:
@@ -57,11 +58,52 @@ def _placed_in(tree) -> list[str]:
     return [args[0]] if isinstance(args[0], str) else []
 
 
+def _milestones(tree, task: Task):
+    """A condition as the milestones it counts: its words, joined as the goal joins them."""
+    name, args = unpack_expression(tree)
+    if name == "and":
+        return ("all", [_milestones(a, task) for a in args])
+    if name == "or":
+        return ("any", [_milestones(a, task) for a in args])
+    if name in ("at_least", "between") and isinstance(args[0], dict) and "tally" in args[0]:
+        _, (parts,) = unpack_expression(args[0])
+        return ("count", [_milestones(p, task) for p in parts], int(args[1]))
+    placed = {n.split("/")[0] for n in _placed_in(tree)}
+    return ("word", concrete_condition(task, tree), placed)
+
+
+def _most(node) -> int:
+    """How many milestones a node counts in all."""
+    if node[0] == "all":
+        return sum(_most(n) for n in node[1])
+    if node[0] == "any":
+        return max((_most(n) for n in node[1]), default=0)
+    if node[0] == "count":
+        return node[2]
+    return 1
+
+
+def _reached(node, state, history, handled) -> int:
+    """How many of a node's milestones hold now, with none of their objects handled."""
+    if node[0] == "all":
+        return sum(_reached(n, state, history, handled) for n in node[1])
+    if node[0] == "any":
+        return max((_reached(n, state, history, handled) for n in node[1]), default=0)
+    if node[0] == "count":
+        return min(sum(_reached(n, state, history, handled) for n in node[1]), node[2])
+    _, condition, placed = node
+    return int(not handled & placed and evaluate(condition, state, history))
+
+
 class Monitor:
     """Evaluates a concrete task's stages and always conditions as an episode runs.
 
     A stage completes once its condition has held for `hold` seconds while none of
     the objects it places is touched by the robot or still moving.
+
+    Progress counts the milestones reached the same way: each word of a stage's
+    condition, with `and` adding, `or` taking the best and `at_least` of a `tally`
+    counting up to its bound. Milestones that hold at the start count from the start.
     """
 
     def __init__(self, task: Task):
@@ -74,6 +116,10 @@ class Monitor:
             name: {n.split("/")[0] for n in _placed_in(stage.condition)}
             for name, stage in self.task.goal.stages.items()
         }
+        self._milestones = {
+            name: _milestones(condition, self.task)
+            for name, condition in self.task.milestones.items()
+        }
 
     def _moved(self, state) -> set[str]:
         """The placed objects that moved since the last check.
@@ -104,6 +150,7 @@ class Monitor:
         self.status = Status()
         self.history = {}
         self._held = dict.fromkeys(self.task.goal.stages, 0.0)
+        self._levels = {name: [math.inf] * _most(tree) for name, tree in self._milestones.items()}
         self._previous = dict.fromkeys(self.task.goal.stages, False)
         self._started = True
         self._finished_at = None
@@ -153,6 +200,7 @@ class Monitor:
             else:
                 self._held[name] = 0.0
             self._previous[name] = eligible and holding
+        self.status.progress = self._progress(state, dt, handled)
         if self.status.succeeded(self.task) and self._finished_at is None:
             self._finished_at = time
         if self.task.time_limit is not None and time > self.task.time_limit:
@@ -160,3 +208,26 @@ class Monitor:
                 self._finished_at is None or self._finished_at > self.task.time_limit
             )
         return self.status
+
+    def _progress(self, state, dt: float, handled: set) -> int:
+        """The milestones reached, each held for its stage's `hold` seconds."""
+        total = 0
+        for name, stage in self.task.goal.stages.items():
+            levels = self._levels[name]
+            if name in self.status.done:
+                total += len(levels)
+                continue
+            eligible = all(p in self.status.done for p in stage.after)
+            tree = self._milestones[name]
+            reached = (
+                _reached(tree, state, self.history, set() if stage.held else handled)
+                if eligible
+                else 0
+            )
+            for level, held in enumerate(levels):
+                if reached <= level:
+                    levels[level] = None
+                else:
+                    levels[level] = 0.0 if held is None else held + dt
+            total += sum(held is not None and held + 1e-9 >= stage.hold for held in levels)
+        return total

{tabletask-0.2.0 → tabletask-0.2.1}/tabletask/language/schema.py

@@ -186,6 +186,7 @@ class Task(msgspec.Struct, omit_defaults=True):
     states: dict[str, dict[str, Any]] = {}
     time_limit: float | None = None
     source: Any = msgspec.field(default=None, name="__source__")
+    milestones: dict[str, Any] = msgspec.field(default_factory=dict, name="__milestones__")
 
     @property
     def root(self) -> Path:

{tabletask-0.2.0 → tabletask-0.2.1}/tabletask/language/transformations.py

@@ -153,6 +153,12 @@ def inline(task: Task) -> Task:
 
 
 def _expand_definitions(task: Task, words: dict) -> Task:
+    return _transform_conditions(task, _definitions(words))
+
+
+def _definitions(words: dict):
+    """A function that expands the given words in a condition."""
+
     def fill(node, bindings):
         if isinstance(node, str):
             if node in bindings:
@@ -201,11 +207,16 @@ def _expand_definitions(task: Task, words: dict) -> Task:
             return {name: [args[0], walk(args[1]), *args[2:]]}
         return {name: [walk(a) for a in args]}
 
-    return _transform_conditions(task, walk)
+    return walk
 
 
 def expand(task: Task) -> Task:
     """Expand `#tag` groups and quantifiers into conditions on named objects."""
+    return _transform_conditions(task, _quantifiers(task))
+
+
+def _quantifiers(task: Task):
+    """A function that expands a task's `#tag` groups and quantifiers in a condition."""
 
     def rename(node, names: dict):
         if isinstance(node, str):
@@ -249,10 +260,27 @@ def expand(task: Task) -> Task:
             return {name: walk(args)}
         return {name: [members(a) if str(a).startswith("#") else walk(a) for a in args]}
 
-    return _transform_conditions(task, walk)
+    return walk
 
 
 def concrete(task: Task) -> Task:
-    """Groups, quantifiers and defined words all expanded, ready to evaluate."""
+    """Groups, quantifiers and defined words all expanded, ready to evaluate.
+
+    The stage conditions as they were before their words were expanded are kept in
+    `milestones`, for counting progress.
+    """
+    grouped = expand(task)
     # expand again after inlining: a defined word's body can hold quantifiers too
-    return expand(inline(expand(task)))
+    result = expand(inline(grouped))
+    if not result.milestones:
+        result.milestones = {name: s.condition for name, s in grouped.goal.stages.items()}
+    return result
+
+
+def concrete_condition(task: Task, condition):
+    """One of a task's conditions expanded as `concrete` expands the task's own."""
+    grow = _quantifiers(task)
+    condition = grow(condition)
+    if task.states:
+        condition = _definitions(task.states)(condition)
+    return grow(_definitions(VOCABULARY)(condition))