refactor: act/obs interface

.gitignore

@@ -13,3 +13,4 @@ run_logs
 .ipynb_checkpoints
 .not*
 dist
+*.parquet

README.md

@@ -161,14 +161,14 @@ Currently located on the branch `diffusion_policy`.
 ## Usage
 To start an vlagents server use the `start-server` command where `kwargs` is a dictionary of the constructor arguments of the policy you want to start e.g.
 ```shell
-# lerobot act (n_action_steps is the executed horizon of the action chunk)
-python -m vlagents start-server lerobot --port 8080 --host 0.0.0.0 --kwargs '{"policy_name": "act", "checkpoint_path": "<path to pretrained_model>", "n_action_steps": 1}'
+# lerobot act
+python -m vlagents start-server lerobot --port 8080 --host 0.0.0.0 --kwargs '{"policy_name": "act", "checkpoint_path": "<path to pretrained_model>"}'
 
 # lerobot pi05
-python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "pi05", "checkpoint_path": "<path to pretrained_model>", "n_action_steps": 1}'
+python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "pi05", "checkpoint_path": "<path to pretrained_model>"}'
 
 # lerobot xvla
-uv run python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "xvla", "checkpoint_path": "<path to pretrained_model>", "n_action_steps": 1, "rename_map": {"head": "image", "left_wrist": "image2", "right_wrist": "image3"}}'
+uv run python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "xvla", "checkpoint_path": "<path to pretrained_model>", "rename_map": {"head": "image", "left_wrist": "image2", "right_wrist": "image3"}}'
 
 
 # octo
@@ -185,23 +185,30 @@ python -m vlagents start-server vjepa --port=20997 --host=0.0.0.0 --kwargs='{"cf
 ```
 
 
+Each policy returns an `Act` action chunk. During evaluation, `EvalEnv.chunk_step` applies the chunk one environment step at a time. Configure `execution_horizon` in an evaluation config to cap how many actions from each chunk are executed before requesting a new one.
+
+Images are resized by `RemoteAgent` before shared-memory or JPEG transport. Set `image_size` in an evaluation config to a `[width, height]` pair (default `[224, 224]`), or `null` to keep native resolution.
+
 There is also the `run-eval-during-training` command to evaluate a model during training, so a single checkpoint.
 The `run-eval-post-training` command evaluates a range of checkpoints in parallel.
 In both cases environment and arguments as well as policy and arguments and wandb config for logging can be passed as CLI arguments.
 
 
 ## Adding your own environment
 ```python
-from vlagents.evaluator_envs import EvaluatorEnv, Obs, Act
+from vlagents import register_env
+from vlagents.envs.interface import EvalEnv
+from vlagents.policies.interface import Act, Obs, SingleAct
 from typing import Any
 
-class YourEnv(EvaluatorEnv):
+class YourEnv(EvalEnv):
+    # Override make_gym() when this environment is not created with gym.make().
 
     def translate_obs(self, obs: dict[str, Any]) -> Obs:
         # translated your observation
         return Obs()
 
-    def step(self, action: Act) -> tuple[Obs, float, bool, bool, dict]:
+    def step(self, action: dict[str, SingleAct]) -> tuple[Obs, float, bool, bool, dict]:
         # step your env
         obs, reward, success, truncated, info = self.env.step(action)
         return self.translate_obs(obs), reward, success, truncated, info
@@ -215,18 +222,18 @@ class YourEnv(EvaluatorEnv):
         # return task instruction
         return "pick up the cube"
 
-    @staticmethod
-    def do_import():
-        # do imports required by your env
+    def do_import(self):
+        # import any packages required by your env
         import libero
 
-EvaluatorEnv.register("your-env-id", YourEnv)
+register_env("your-env-id", YourEnv)
 ```
 
 ## Adding your own policy
 ```python
-from vlagents.policies import Agent, AGENTS
-from vlagents.evaluator_envs import Obs, Act
+from vlagents import register_agent
+from vlagents.policies.interface import Agent
+from vlagents.policies.interface import Obs, Act
 from typing import Any
 import numpy as np
 
@@ -245,20 +252,20 @@ class YourAgent(Agent):
 
     def close(self, *args, **kwargs):
         pass
-AGENTS["your-agent-id"] = YourAgent
+register_agent("your-agent-id", YourAgent)
 ```
 
 
 
 ## Contribution
 
 ### New Policy
-In order to extend the library with a new policy network, extend the `Agent` class in [policies.py](src/vlagents/policies.py).
+In order to extend the library with a new policy network, extend the `Agent` class in [policies/interface.py](src/vlagents/policies/interface.py).
 It is important to only invoke policy specific imports in the class functions, as each policy can have its own dependencies.
 
 
 ### New Environment
-In order to extend the library with a new agent environment, extend the `EvaluatorEnv` class in [evaluator_envs.py](src/vlagents/evaluator_envs.py).
+In order to extend the library with a new agent environment, extend the `EvalEnv` class in [envs/interface.py](src/vlagents/envs/interface.py).
 
 
 ### Developer Tools

pyproject.toml

@@ -103,6 +103,8 @@ target-version = ["py310"]
 
 [tool.isort]
 profile = "black"
+known_first_party = ["vlagents"]
+known_third_party = ["lerobot", "openpi"]
 
 [tool.mypy]
 ignore_missing_imports = true
@@ -122,3 +124,6 @@ version_scheme = "pep440"
 version_provider = "pep621"
 update_changelog_on_bump = true
 major_version_zero = true
+version_files = [
+  "src/vlagents/__init__.py:__version__",
+]

src/tests/test_connection.py

@@ -4,47 +4,67 @@
 import numpy as np
 
 from vlagents.client import RemoteAgent
-from vlagents.evaluator_envs import start_server
-from vlagents.policies import Act, Obs
+from vlagents.eval import start_server
+from vlagents.policies.interface import Obs, SingleObs
+
+
+def _make_obs(data: np.ndarray, instruction: str = "do something") -> Obs:
+    return Obs(
+        obs={"right": SingleObs(cameras={"rgb_side": data})},
+        language_instruction=instruction,
+    )
+
+
+def _single_robot_action_info(act):
+    step = act.acts[0]
+    robot_action = step["right"]
+    return robot_action.action, robot_action.gripper, robot_action.done, robot_action.info
 
 
 def _test_connection(agent: RemoteAgent):
     data = np.zeros((256, 256, 3), dtype=np.uint8)
     data[2, 0, 0] = 16
-    obs = Obs(cameras=dict(rgb_side=data))
-    instruction = "do something"
-    reset_info = agent.reset(obs, instruction)
-    assert reset_info["instruction"] == instruction
-    assert reset_info["shapes"] == {"rgb_side": [256, 256, 3]}
-    assert reset_info["dtype"] == {"rgb_side": "uint8"}
-    assert (reset_info["data"]["rgb_side"] == data).all()
+
+    first = agent.act(_make_obs(data))
+    first_action, first_gripper, first_done, first_info = _single_robot_action_info(first)
+    assert first_info["shapes"] == {"rgb_side": [224, 224, 3]}
+    assert first_info["dtype"] == {"rgb_side": "uint8"}
+    assert first_info["data"]["rgb_side"].shape == (224, 224, 3)
+    assert np.all(first_action == np.array([0, 0, 0, 0, 0, 0], dtype=np.float32))
+    assert first_gripper == 0.0
+    assert not first_done
 
     data[0, 0, 2] = 1
-    a1 = agent.act(Obs(cameras=dict(rgb_side=data)))
-    assert a1.info["shapes"] == {"rgb_side": [256, 256, 3]}
-    assert a1.info["dtype"] == {"rgb_side": "uint8"}
-    assert (a1.info["data"]["rgb_side"] == data).all()
-    assert np.all(a1.action == np.array([0, 0, 0, 0, 0, 0, 0], dtype=np.float32))
-    assert not a1.done
-
-    data[0, 2, 0] = 1
-    a1 = agent.act(Obs(cameras=dict(rgb_side=data)))
-    assert a1.info["shapes"] == {"rgb_side": [256, 256, 3]}
-    assert a1.info["dtype"] == {"rgb_side": "uint8"}
-    assert (a1.info["data"]["rgb_side"] == data).all()
-    assert np.all(a1.action == np.array([0, 0, 0, 0, 0, 0, 1], dtype=np.float32))
-    assert not a1.done
+    second = agent.act(_make_obs(data))
+    second_action, second_gripper, second_done, second_info = _single_robot_action_info(second)
+    assert second_info["shapes"] == {"rgb_side": [224, 224, 3]}
+    assert second_info["dtype"] == {"rgb_side": "uint8"}
+    assert second_info["data"]["rgb_side"].shape == (224, 224, 3)
+    assert np.all(second_action == np.array([0, 0, 0, 0, 0, 0], dtype=np.float32))
+    assert second_gripper == 1.0
+    assert not second_done
 
 
 def _test_connection_jpeg(agent: RemoteAgent):
     data = np.zeros((256, 256, 3), dtype=np.uint8)
-    obs = Obs(cameras=dict(rgb_side=data))
-    instruction = "do something"
-    reset_info = agent.reset(obs, instruction)
-    assert reset_info["instruction"] == instruction
-    assert reset_info["shapes"] == {"rgb_side": [256, 256, 3]}
-    assert reset_info["dtype"] == {"rgb_side": "uint8"}
-    assert (reset_info["data"]["rgb_side"] == data).all()
+    act = agent.act(_make_obs(data))
+    action, gripper, done, info = _single_robot_action_info(act)
+    assert info["shapes"] == {"rgb_side": [224, 224, 3]}
+    assert info["dtype"] == {"rgb_side": "uint8"}
+    assert info["data"]["rgb_side"].shape == (224, 224, 3)
+    assert np.all(action == np.array([0, 0, 0, 0, 0, 0], dtype=np.float32))
+    assert gripper == 0.0
+    assert not done
+
+
+def _test_connection_without_resize(agent: RemoteAgent):
+    data = np.zeros((256, 256, 3), dtype=np.uint8)
+    data[2, 0, 0] = 16
+    act = agent.act(_make_obs(data))
+    _, _, _, info = _single_robot_action_info(act)
+    assert info["shapes"] == {"rgb_side": [256, 256, 3]}
+    assert info["dtype"] == {"rgb_side": "uint8"}
+    np.testing.assert_array_equal(info["data"]["rgb_side"], data)
 
 
 def test_connection_numpy_serialization():
@@ -78,3 +98,14 @@ def test_connection_numpy_jpeg():
                 sleep(0.1)
             _test_connection_jpeg(agent)
         p.send_signal(subprocess.signal.SIGINT)
+
+
+def test_connection_preserves_native_resolution():
+    with start_server("test", {}, 8080, "localhost") as p:
+        sleep(2)
+        agent = RemoteAgent("localhost", 8080, "test", image_size=None)
+        with agent:
+            while not agent.is_initialized():
+                sleep(0.1)
+            _test_connection_without_resize(agent)
+        p.send_signal(subprocess.signal.SIGINT)

src/tests/test_eval_merge.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from vlagents.__main__ import _merge_env_split_results
-from vlagents.evaluator_envs import EvalConfig
+from vlagents.envs.interface import EvalConfig
 
 
 def test_merge_env_split_results_keeps_distinct_seeded_cfgs():

src/tests/test_libero.py

@@ -1,13 +1,8 @@
 if __name__ == "__main__":
-    import datetime
     import os
 
-    import numpy as np
-    from lerobot.envs.libero import LiberoEnv
-    from PIL import Image
-
     from vlagents.__main__ import _run_eval
-    from vlagents.evaluator_envs import AgentConfig, EvalConfig
+    from vlagents.envs.interface import AgentConfig, EvalConfig
 
     # main_app()
     # test

src/vlagents/__init__.py

@@ -1 +1,33 @@
-__all__ = ["__doc__", "__version__", "policies", "client", "server"]
+AGENTS = {}
+ENVS = {}
+
+
+def register_agent(name: str, agent_class: type["Agent"]) -> None:
+    """
+    Register an agent class with a given name.
+
+    Args:
+        name (str): The name of the agent.
+        agent_class (type[Agent]): The agent class to register.
+    """
+    AGENTS[name] = agent_class
+
+
+def register_env(name: str, env_class: type["EvalEnv"]) -> None:
+    """
+    Register an environment class with a given name.
+
+    Args:
+        name (str): The name of the environment.
+        env_class (type[EvalEnv]): The environment class to register.
+    """
+    ENVS[name] = env_class
+
+
+from vlagents.envs import duobench, libero, maniskill  # noqa: E402, F401
+from vlagents.envs.interface import EvalEnv
+from vlagents.policies import lerobot, octo, openpi, openvla, vjepa  # noqa: E402, F401
+from vlagents.policies.interface import Agent
+
+__version__ = "0.2.0"
+__all__ = ["__doc__", "__version__", "AGENTS", "ENVS", "register_agent", "register_env", "EvalEnv", "Agent"]

src/vlagents/__main__.py

@@ -15,8 +15,9 @@
 # when started from jupyter notebook
 os.environ["MPLBACKEND"] = "Agg"
 
-from vlagents.evaluator_envs import AgentConfig, EvalConfig, evaluation, write_results
-from vlagents.policies import AGENTS
+from vlagents import AGENTS
+from vlagents.envs.interface import AgentConfig, EvalConfig
+from vlagents.eval import evaluation, write_results
 from vlagents.server import AgentService
 
 main_app = typer.Typer(help="CLI tool for the vlagents library.")
@@ -207,7 +208,7 @@ def run_eval(
     n_gpus: Annotated[int, typer.Option(help="Number of gpus to run.")] = 1,
     eval_cfgs: Annotated[
         str, typer.Option(help="Evaluation configurations.")
-    ] = '[{"env": "rcs/SimplePickUpSim-v0", "kwargs": {}}]',
+    ] = '[{"env_id": "rcs/SimplePickUpSim-v0", "env_kwargs": {}}]',
     agent_cfg: Annotated[
         str, typer.Option(help="Agent configuration.")
     ] = '{"host": "localhost", "port": 8080, "agent_name": "Test", "agent_kwargs": {}, "python_path": "python"}',