Add data folder with demo_types and collect method + tests

Author: zahraabashir

src/programmatic_policy_learning/data/collect.py

@@ -0,0 +1,47 @@
+"""Demo collection utilities."""
+
+import logging
+from typing import Any, Callable
+
+import numpy as np
+
+from programmatic_policy_learning.data.demo_types import Demo, Trajectory
+
+EnvFactory = Callable[[], Any]
+
+
+def collect_demo(
+    env_factory: EnvFactory, expert: Any, max_demo_length: int | float = np.inf
+) -> Trajectory:
+    """Collect a demonstration trajectory from an environment using an expert
+    policy."""
+
+    env = env_factory()
+    reset_out = env.reset()
+    if isinstance(reset_out, tuple) and len(reset_out) == 2:
+        obs, info = reset_out
+    else:
+        obs = reset_out
+        info = {}
+        logging.warning("env.reset() returned a single value (old gym API)")
+
+    steps: list[Demo] = []
+    t = 0
+    expert.reset(obs, info)
+    while True:
+        action = expert.step()
+        steps.append(Demo(obs=obs, act=action))
+        step_out = env.step(action)
+        # handle gym vs. gymnasium
+        if len(step_out) == 4:
+            obs, reward, done, _ = step_out
+            terminated, truncated = done, False
+        else:
+            obs, reward, terminated, truncated, _ = step_out
+        t += 1
+        if terminated or truncated or (t >= max_demo_length):
+            if not reward > 0:
+                # keep behavior parity with original: warn if didn’t succeed
+                print("WARNING: demo did not succeed!")
+            break
+    return Trajectory(steps=steps)

src/programmatic_policy_learning/data/demo_types.py

@@ -0,0 +1,22 @@
+"""Types for demonstrations and trajectories."""
+
+from dataclasses import dataclass
+from typing import Generic, Sequence, TypeVar
+
+ObsT = TypeVar("ObsT")
+ActT = TypeVar("ActT")
+
+
+@dataclass(frozen=True)
+class Demo(Generic[ObsT, ActT]):
+    """A single demonstration step, containing observation and action."""
+
+    obs: ObsT
+    act: ActT
+
+
+@dataclass(frozen=True)
+class Trajectory(Generic[ObsT, ActT]):
+    """A sequence of demonstration steps (Demo), called a trajectory."""
+
+    steps: Sequence[Demo[ObsT, ActT]]

tests/data/test_collect.py

@@ -0,0 +1,79 @@
+"""Tests for data/collect.py."""
+
+from typing import Any
+
+import numpy as np
+from omegaconf import DictConfig, OmegaConf
+
+from programmatic_policy_learning.approaches.random_actions import RandomActionsApproach
+from programmatic_policy_learning.data.collect import collect_demo
+from programmatic_policy_learning.data.demo_types import Demo, Trajectory
+from programmatic_policy_learning.envs.registry import EnvRegistry
+
+
+def test_collect_demo_returns_trajectory_DummyEnv():
+    """Test that collect_demo returns a Trajectory with Demo steps,
+    DummyEnv."""
+
+    class DummySpace:
+        """DummySpace."""
+
+        def sample(self):
+            """Sample."""
+            return 0
+
+        def seed(self, _):
+            """Seed."""
+            pass  # pylint: disable=unnecessary-pass
+
+    class DummyEnv:
+        """DummyEnv."""
+
+        def reset(self):
+            """Reset."""
+            return np.zeros((2, 2)), {}
+
+        def step(self, _):
+            """Step."""
+            return np.zeros((2, 2)), 1.0, True, False, {}
+
+        @property
+        def observation_space(self):
+            """Observation space."""
+            return DummySpace()
+
+        @property
+        def action_space(self):
+            """Action space."""
+            return DummySpace()
+
+    env = DummyEnv()
+    env_factory = lambda: env  # returns a new environment each time you call
+
+    expert = RandomActionsApproach(
+        "TEST", env.observation_space, env.action_space, seed=1
+    )
+    traj = collect_demo(env_factory, expert, max_demo_length=5)
+    assert isinstance(traj, Trajectory)
+    assert all(isinstance(step, Demo) for step in traj.steps)
+    assert len(traj.steps) > 0
+
+
+def test_collect_demo_with_real_env():
+    """Test collect_demo with a real environment using EnvRegistry."""
+    cfg: DictConfig = OmegaConf.create(
+        {
+            "provider": "ggg",
+            "make_kwargs": {"id": "TwoPileNim0-v0"},
+        }
+    )
+    registry = EnvRegistry()
+    env_factory = lambda: registry.load(cfg)
+    env: Any = env_factory()  # type: ignore
+    expert = RandomActionsApproach(
+        "TEST", env.observation_space, env.action_space, seed=1
+    )
+    traj = collect_demo(env_factory, expert, max_demo_length=5)
+    assert isinstance(traj, Trajectory)
+    assert all(isinstance(step, Demo) for step in traj.steps)
+    assert len(traj.steps) > 0