diff --git a/.gitignore b/.gitignore index ffad55d..202daf5 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,4 @@ run_logs .ipynb_checkpoints .not* dist +*.parquet diff --git a/README.md b/README.md index 75e2e8c..6a24914 100644 --- a/README.md +++ b/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": "", "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": ""}' # lerobot pi05 -python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "pi05", "checkpoint_path": "", "n_action_steps": 1}' +python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "pi05", "checkpoint_path": ""}' # lerobot xvla -uv run python -m vlagents start-server lerobot --port 20000 --host 0.0.0.0 --kwargs '{"policy_name": "xvla", "checkpoint_path": "", "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": "", "rename_map": {"head": "image", "left_wrist": "image2", "right_wrist": "image3"}}' # octo @@ -185,6 +185,10 @@ 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. @@ -192,16 +196,19 @@ In both cases environment and arguments as well as policy and arguments and wand ## 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,7 +252,7 @@ class YourAgent(Agent): def close(self, *args, **kwargs): pass -AGENTS["your-agent-id"] = YourAgent +register_agent("your-agent-id", YourAgent) ``` @@ -253,12 +260,12 @@ AGENTS["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 diff --git a/pyproject.toml b/pyproject.toml index 0dc4223..8cf3355 100644 --- a/pyproject.toml +++ b/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__", +] diff --git a/src/tests/test_connection.py b/src/tests/test_connection.py index 7212057..0c82bf2 100644 --- a/src/tests/test_connection.py +++ b/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) diff --git a/src/tests/test_eval_merge.py b/src/tests/test_eval_merge.py index cae3697..999a09e 100644 --- a/src/tests/test_eval_merge.py +++ b/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(): diff --git a/src/tests/test_libero.py b/src/tests/test_libero.py index b37cc53..9602890 100644 --- a/src/tests/test_libero.py +++ b/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 diff --git a/src/vlagents/__init__.py b/src/vlagents/__init__.py index 0ad55f6..a092803 100644 --- a/src/vlagents/__init__.py +++ b/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"] diff --git a/src/vlagents/__main__.py b/src/vlagents/__main__.py index 5176894..f0e8aa7 100644 --- a/src/vlagents/__main__.py +++ b/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"}', diff --git a/src/vlagents/client.py b/src/vlagents/client.py index 4e4f670..6cae5b7 100644 --- a/src/vlagents/client.py +++ b/src/vlagents/client.py @@ -2,27 +2,51 @@ import dataclasses from dataclasses import asdict from multiprocessing import shared_memory -from typing import Any +from typing import get_args, get_origin import json_numpy import numpy as np import rpyc import simplejpeg - -from vlagents.policies import Act, Agent, CameraDataType, Obs, SharedMemoryPayload - - -def dataclass_from_dict(klass, d): - # https://stackoverflow.com/questions/53376099/python-dataclass-from-a-nested-dict - try: +from PIL import Image + +from vlagents.policies.interface import ( + Act, + Agent, + CameraDataType, + Obs, + SharedMemoryPayload, +) + + +def dataclass_from_dict(klass, value): + origin = get_origin(klass) + if origin is dict: + key_type, value_type = get_args(klass) + return { + dataclass_from_dict(key_type, key): dataclass_from_dict(value_type, item) for key, item in value.items() + } + if origin is list: + (item_type,) = get_args(klass) + return [dataclass_from_dict(item_type, item) for item in value] + + if dataclasses.is_dataclass(klass): fieldtypes = {f.name: f.type for f in dataclasses.fields(klass)} - return klass(**{f: dataclass_from_dict(fieldtypes[f], d[f]) for f in d}) - except: - return d # Not a dataclass field + return klass(**{field: dataclass_from_dict(fieldtypes[field], value[field]) for field in value}) + + return value class RemoteAgent(Agent): - def __init__(self, host: str, port: int, model: str, on_same_machine: bool = False, jpeg_encoding: bool = False): + def __init__( + self, + host: str, + port: int, + model: str, + on_same_machine: bool = False, + jpeg_encoding: bool = False, + image_size: tuple[int, int] | None = (224, 224), + ): """Connect to a remote agent service. Args: @@ -33,12 +57,15 @@ def __init__(self, host: str, port: int, model: str, on_same_machine: bool = Fal shared memory for more efficient communication. Defaults to False. jpeg_encoding (bool, optional): If True the image data is jpeg encoded for smaller transfer size. Defaults to False. + image_size (tuple[int, int] | None, optional): Image size as (width, height) applied before + serialization. Set to None to retain native resolution. Defaults to (224, 224). """ self.host = host self.port = port self.model = model self.on_same_machine = on_same_machine self.jpeg_encoding = jpeg_encoding + self.image_size = self._validate_image_size(image_size) self._shm: dict[str, shared_memory.SharedMemory] = {} self.c = None self._connect() @@ -58,6 +85,7 @@ def reconnect( model: str | None = None, on_same_machine: bool | None = None, jpeg_encoding: bool | None = None, + image_size: tuple[int, int] | None = None, ): if self.c is not None: try: @@ -72,6 +100,8 @@ def reconnect( self.model = model if on_same_machine is not None: self.on_same_machine = on_same_machine + if image_size is not None: + self.image_size = self._validate_image_size(image_size) if jpeg_encoding is not None: self.jpeg_encoding = jpeg_encoding self._connect() @@ -83,33 +113,75 @@ def ensure_connected(self): except Exception: self.reconnect() + def _to_shared_memory_payload(self, shm_key: str, image: np.ndarray) -> SharedMemoryPayload: + if shm_key not in self._shm or self._shm[shm_key].size < image.nbytes: + if shm_key in self._shm: + self._shm[shm_key].close() + self._shm[shm_key].unlink() + self._shm[shm_key] = shared_memory.SharedMemory(create=True, size=image.nbytes) + image_shared = np.ndarray(image.shape, buffer=self._shm[shm_key].buf, dtype=image.dtype) + image_shared[:] = image[:] + return SharedMemoryPayload( + shm_name=self._shm[shm_key].name, + shape=image.shape, + dtype=image.dtype.name, + ) + + @staticmethod + def _to_jpeg_payload(image: np.ndarray) -> str: + return base64.urlsafe_b64encode(simplejpeg.encode_jpeg(np.ascontiguousarray(image))).decode("utf-8") + + @staticmethod + def _validate_image_size(image_size: tuple[int, int] | None) -> tuple[int, int] | None: + if image_size is None: + return None + if len(image_size) != 2 or any(not isinstance(size, (int, np.integer)) or size <= 0 for size in image_size): + message = "image_size must be a (width, height) pair of positive integers or None" + raise ValueError(message) + return tuple(int(size) for size in image_size) + + def _resize_image(self, image: np.ndarray) -> np.ndarray: + if self.image_size is None or image.shape[:2] == self.image_size[::-1]: + return image + if image.ndim == 3: + return np.asarray(Image.fromarray(image).resize(self.image_size, Image.Resampling.BILINEAR)) + if image.ndim == 4: + return np.stack([self._resize_image(frame) for frame in image]) + message = f"Expected an HWC image or NHWC batch, got shape {image.shape}" + raise ValueError(message) + def _process(self, obs: Obs) -> Obs: - if self.on_same_machine: - camera_dict = {} - for camera_name, camera_data in obs.cameras.items(): - assert isinstance(camera_data, np.ndarray) - if camera_name not in self._shm: - self._shm[camera_name] = shared_memory.SharedMemory(create=True, size=camera_data.nbytes) - camera_shared = np.ndarray( - camera_data.shape, buffer=self._shm[camera_name].buf, dtype=camera_data.dtype - ) - camera_shared[:] = camera_data[:] - camera_dict[camera_name] = SharedMemoryPayload( - shm_name=self._shm[camera_name].name, - shape=camera_data.shape, - dtype=camera_data.dtype.name, - ) - obs.cameras = camera_dict - obs.camera_data_type = CameraDataType.SHARED_MEMORY - elif self.jpeg_encoding: - camera_dict = {} - for camera_name, camera_data in obs.cameras.items(): - assert isinstance(camera_data, np.ndarray) - camera_dict[camera_name] = base64.urlsafe_b64encode( - simplejpeg.encode_jpeg(np.ascontiguousarray(camera_data)) - ).decode("utf-8") - obs.cameras = camera_dict - obs.camera_data_type = CameraDataType.JPEG_ENCODED + for robot_name, single_obs in obs.obs.items(): + single_obs.cameras = { + camera_name: self._resize_image(camera_data) for camera_name, camera_data in single_obs.cameras.items() + } + if self.on_same_machine: + camera_dict = {} + for camera_name, camera_data in single_obs.cameras.items(): + assert isinstance(camera_data, np.ndarray) + camera_dict[camera_name] = self._to_shared_memory_payload( + f"{robot_name}:{camera_name}", camera_data + ) + single_obs.cameras = camera_dict + single_obs.camera_data_type = CameraDataType.SHARED_MEMORY + elif self.jpeg_encoding: + camera_dict = {} + for camera_name, camera_data in single_obs.cameras.items(): + assert isinstance(camera_data, np.ndarray) + camera_dict[camera_name] = self._to_jpeg_payload(camera_data) + single_obs.cameras = camera_dict + single_obs.camera_data_type = CameraDataType.JPEG_ENCODED + + if obs.goal_image is not None: + assert isinstance(obs.goal_image, np.ndarray) + obs.goal_image = self._resize_image(obs.goal_image) + if self.on_same_machine: + obs.goal_image = self._to_shared_memory_payload("goal_image", obs.goal_image) + obs.goal_image_data_type = CameraDataType.SHARED_MEMORY + elif self.jpeg_encoding: + obs.goal_image = self._to_jpeg_payload(obs.goal_image) + obs.goal_image_data_type = CameraDataType.JPEG_ENCODED + return obs def act(self, obs: Obs) -> Act: @@ -124,18 +196,6 @@ def act(self, obs: Obs) -> Act: assert self.c is not None return dataclass_from_dict(Act, json_numpy.loads(self.c.root.act(obs))) - def reset(self, obs: Obs, instruction: Any, **kwargs) -> dict[str, Any]: - obs = self._process(obs) - obs_dict = asdict(obs) - # info - try: - assert self.c is not None - return json_numpy.loads(self.c.root.reset(json_numpy.dumps((obs_dict, instruction, kwargs)))) - except Exception: - self.reconnect() - assert self.c is not None - return json_numpy.loads(self.c.root.reset(json_numpy.dumps((obs_dict, instruction, kwargs)))) - def git_status(self) -> str: assert self.c is not None return json_numpy.loads(self.c.root.git_status()) @@ -155,9 +215,12 @@ def close(self): if __name__ == "__main__": # to test the connection + from vlagents.policies.interface import SingleObs + agent = RemoteAgent("localhost", 8080, "test") - obs = Obs(cameras={"rgb_side": np.zeros((256, 256, 3), dtype=np.uint8)}) - instruction = "do something" - agent.reset(obs, instruction) + obs = Obs( + obs={"right": SingleObs(cameras={"rgb_side": np.zeros((256, 256, 3), dtype=np.uint8)})}, + language_instruction="do something", + ) print(agent.act(obs)) print(agent.act(obs)) diff --git a/src/vlagents/envs/__init__.py b/src/vlagents/envs/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/vlagents/envs/duobench.py b/src/vlagents/envs/duobench.py new file mode 100644 index 0000000..15fcff5 --- /dev/null +++ b/src/vlagents/envs/duobench.py @@ -0,0 +1,91 @@ +from typing import Any, ClassVar + +import numpy as np + +from vlagents import register_env +from vlagents.envs.interface import EvalEnv +from vlagents.policies.interface import Obs, SingleAct, SingleObs + + +class RCSDuoBench(EvalEnv): + INSTRUCTIONS: ClassVar[dict[str, str]] = {} + + def __init__(self, env_id, **env_kwargs): + self.robot_keys: str = env_kwargs.pop("robot_keys", ["left", "right"]) + self.control_mode: str = env_kwargs.pop("control_mode", "joints") + self._instruction: str | None = None + super().__init__(env_id, **env_kwargs) + + def translate_obs(self, obs: dict[str, Any]) -> Obs: + cameras = {key: obs["frames"][key]["rgb"]["data"] for key in obs["frames"]} + return Obs( + obs={ + robot_key: SingleObs( + cameras=cameras.copy(), + joints=np.asarray(obs[robot_key]["joints"], dtype=np.float32), + gripper=float(obs[robot_key]["gripper"]), + ) + for robot_key in self.robot_keys + }, + language_instruction=self.language_instruction, + ) + + def step(self, action: dict[str, SingleAct]) -> tuple[Obs, float, bool, bool, dict]: + env_action = {} + for robot in self.robot_keys: + robot_action = action[robot] + gripper = 0.0 if robot_action.gripper is None else robot_action.gripper + if self.control_mode == "joints": + env_action[robot] = { + "joints": np.asarray(robot_action.action, dtype=np.float32), + "gripper": np.asarray([gripper], dtype=np.float32), + } + else: + env_action[robot] = { + "xyzrpy": np.asarray(robot_action.action, dtype=np.float32), + "gripper": np.asarray([gripper], dtype=np.float32), + } + obs, reward, success, truncated, info = self.env.step(env_action) + r = float(reward) + + return self.translate_obs(obs), r, success, truncated, info + + def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Obs, dict[str, Any]]: + obs, info = self.env.reset(seed=seed, options=options) + self._instruction = info["instruction"] + return self.translate_obs(obs), info + + @property + def language_instruction(self) -> str: + assert self._instruction is not None + return self._instruction + + @staticmethod + def do_import(): + import rcs + from rcs_duobench.tasks import ( + ball_maze, + bin_sort, + block_balance, + carry_pot, + hinge_chest, + join_blocks, + pour_marbles, + spring_door, + transfer_cube, + transfer_gate, + transfer_reorient, + ) + + +register_env("duobench/ball_maze", RCSDuoBench) +register_env("duobench/bin_sort", RCSDuoBench) +register_env("duobench/block_balance", RCSDuoBench) +register_env("duobench/carry_pot", RCSDuoBench) +register_env("duobench/join_blocks", RCSDuoBench) +register_env("duobench/hinge_chest", RCSDuoBench) +register_env("duobench/pour_marbles", RCSDuoBench) +register_env("duobench/spring_door", RCSDuoBench) +register_env("duobench/transfer_cube", RCSDuoBench) +register_env("duobench/transfer_gate", RCSDuoBench) +register_env("duobench/transfer_reorient", RCSDuoBench) diff --git a/src/vlagents/envs/interface.py b/src/vlagents/envs/interface.py new file mode 100644 index 0000000..2708ebd --- /dev/null +++ b/src/vlagents/envs/interface.py @@ -0,0 +1,91 @@ +import logging +from abc import ABC +from dataclasses import dataclass +from typing import Any + +import gymnasium as gym + +from vlagents import ENVS +from vlagents.policies.interface import Act, Obs, SingleAct + +logging.basicConfig( + format="%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + level=logging.INFO, +) + + +class EvalEnv(ABC): + + def __init__(self, env_id: str, execution_horizon: int | None = None, **env_kwargs) -> None: + self.env_id = env_id + self.env_kwargs = env_kwargs + self.execution_horizon = execution_horizon + self.do_import() + self.env = self.make_gym() + self.last_chunk_steps = 0 + + def chunk_step(self, actions: Act, max_steps: int | None = None) -> tuple[Obs, float, bool, bool, dict[str, Any]]: + if not actions.acts: + raise ValueError("Agents must return at least one action") + rewards = [] + self.last_chunk_steps = 0 + for action in actions.acts: + if max_steps is not None and self.last_chunk_steps >= max_steps: + break + obs, reward, done, truncated, info = self.step(action) + rewards.append(reward) + self.last_chunk_steps += 1 + if ( + done + or truncated + or (self.execution_horizon is not None and self.last_chunk_steps >= self.execution_horizon) + ): + break + if not rewards: + raise ValueError("max_steps must allow at least one environment step") + return obs, sum(rewards), done, truncated, info + + def step(self, action: dict[str, SingleAct]) -> tuple[Obs, float, bool, bool, dict]: + raise NotImplementedError + + def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Obs, dict[str, Any]]: + raise NotImplementedError + + @property + def language_instruction(self) -> str: + raise NotImplementedError + + def make_gym(self) -> gym.Env: + return gym.make(self.env_id, **self.env_kwargs) + + def do_import(self): + raise NotImplementedError + + @staticmethod + def from_id(env_id: str, execution_horizon: int | None = None, **env_kwargs) -> "EvalEnv": + if env_id not in ENVS: + raise ValueError(f"Unknown environment id {env_id}. Available environments: {list(ENVS.keys())}") + return ENVS[env_id](env_id, execution_horizon=execution_horizon, **env_kwargs) + + +@dataclass +class EvalConfig: + env_id: str + env_kwargs: dict[str, Any] + execution_horizon: int | None = None + max_steps_per_episode: int = 100 + seed: int = 42 + same_machine: bool = False + jpeg_encoding: bool = False + image_size: tuple[int, int] | None = (224, 224) + + +@dataclass +class AgentConfig: + host: str + agent_name: str + agent_kwargs: dict[str, Any] + python_path: str = "python" + """modify this if you want to use a specific python environment """ + port: int = 8080 diff --git a/src/vlagents/envs/libero.py b/src/vlagents/envs/libero.py new file mode 100644 index 0000000..9bf0e0f --- /dev/null +++ b/src/vlagents/envs/libero.py @@ -0,0 +1,146 @@ +import logging +import os +from typing import Any + +import numpy as np + +from vlagents import register_env +from vlagents.envs.interface import EvalEnv +from vlagents.policies.interface import Obs, SingleAct, SingleObs + + +class Libero(EvalEnv): + def __init__(self, env_id: str, reset_steps: int = 14, execution_horizon: int | None = None, **env_kwargs) -> None: + """ + For supported env_kwargs checkout ControlEnv class in libero. + We add the following env_kwargs on top: + - task_id (int): libero task id for given task suite. The number of tasks per task suite can checked with Libero.n_tasks(env_id). Defaults to 0. + - control_mode (str): either 'relative' or 'absolute'. Defaults to 'relative'. + + """ + logging.info("Creating Libero env") + self.reset_steps = reset_steps + self.control_mode = env_kwargs.pop("control_mode", "relative") + super().__init__(env_id, execution_horizon=execution_horizon, **env_kwargs) + + @staticmethod + def n_tasks(env_id: str) -> int: + from libero.libero import benchmark + + benchmark_dict = benchmark.get_benchmark_dict() + task_suite = benchmark_dict[env_id]() + return task_suite.n_tasks + + @staticmethod + def _make_gym(env_id, **env_kwargs): + from libero.libero import benchmark, get_libero_path + from libero.libero.envs import OffScreenRenderEnv + + benchmark_dict = benchmark.get_benchmark_dict() + + task_suite = benchmark_dict[env_id]() + task_id = min(max(env_kwargs.pop("task_id", 0), 0), task_suite.n_tasks - 1) + task = task_suite.get_task(task_id) + + task_bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file) + env = OffScreenRenderEnv( + bddl_file_name=task_bddl_file, + **env_kwargs, + ) + + return env, task.language, task.name, task_suite, task_id, task + + def make_gym(self): + ( + env, + self._language_instruction, + self.task_name, + self.task_suite, + self.task_id, + self.task, + ) = self._make_gym(self.env_id, **self.env_kwargs) + logging.info( + f"Created Libero env, task suite: {self.env_id}, task id: {self.task_id}, task name {self.task_name}, instruction: {self._language_instruction}" + ) + return env + + def do_import(self): + # _make_gym imports LIBERO lazily so importing vlagents does not require it. + pass + + def translate_obs(self, obs: dict[str, Any]) -> Obs: + joints = None + if "robot0_joint_pos" in obs: + joints = np.asarray(obs["robot0_joint_pos"], dtype=np.float32) + gripper = float(np.asarray(obs["robot0_gripper_qpos"]).reshape(-1)[0] / 0.04) + return Obs( + obs={ + "robot0": SingleObs( + cameras={ + "rgb_side": obs["agentview_image"][::-1], + "rgb_wrist": obs["robot0_eye_in_hand_image"][::-1], + }, + joints=joints, + gripper=gripper, + ) + }, + language_instruction=self.language_instruction, + ) + + def step(self, action: dict[str, SingleAct]) -> tuple[Obs, float, bool, bool, dict]: + # change gripper to libero format (-1, 1) where -1 is open + assert len(action) == 1, "Libero expects a single robot action" + _, robot_action = next(iter(action.items())) + if robot_action.gripper is None: + raise ValueError("Libero expects a gripper value in SingleAct.gripper") + act = np.concatenate( + [ + np.asarray(robot_action.action, dtype=np.float32), + np.asarray([robot_action.gripper], dtype=np.float32), + ] + ) + act[-1] = (1 - act[-1]) * 2 - 1.0 + obs, reward, done, info = self.env.step(act) + success = self.env.check_success() + return self.translate_obs(obs), reward, success, done, info + + def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Obs, dict[str, Any]]: + if seed is not None: + self.env.seed(seed) + obs = self.env.reset() + init_states = self.task_suite.get_task_init_states( + self.task_id + ) # for benchmarking purpose, we fix the a set of initial states + init_state_id = 0 + self.env.set_init_state(init_states[init_state_id]) + + for robot in self.env.robots: + robot.controller.use_delta = True + for _ in range(self.reset_steps): + # steps the environment to filter out falling objects + obs, _, _, _ = self.env.step( + np.zeros(8) if "JOINT" in self.env_kwargs.get("controller", "OSC_POSE") else np.zeros(7) + ) + + if self.control_mode == "absolute": + for robot in self.env.robots: + robot.controller.use_delta = False + elif self.control_mode == "relative": + for robot in self.env.robots: + robot.controller.use_delta = True + else: + raise ValueError(f"Invalid control mode: {self.control_mode}, use 'absolute' or 'relative'.") + + return self.translate_obs(obs), {} + + @property + def language_instruction(self) -> str: + return self._language_instruction + + +register_env("libero_10", Libero) +register_env("libero_90", Libero) +register_env("libero_100", Libero) +register_env("libero_spatial", Libero) +register_env("libero_object", Libero) +register_env("libero_goal", Libero) diff --git a/src/vlagents/envs/maniskill.py b/src/vlagents/envs/maniskill.py new file mode 100644 index 0000000..9e92175 --- /dev/null +++ b/src/vlagents/envs/maniskill.py @@ -0,0 +1,131 @@ +import copy +import logging +from typing import Any, ClassVar + +import gymnasium as gym +import numpy as np + +from vlagents import register_env +from vlagents.envs.interface import EvalEnv +from vlagents.policies.interface import Obs, SingleAct, SingleObs + + +class HumanCameraWrapper(gym.ObservationWrapper): + """ + Flattens the rgbd mode observations into a dictionary with two keys, "rgbd" and "state" + + Args: + rgb (bool): Whether to include rgb images in the observation + depth (bool): Whether to include depth images in the observation + state (bool): Whether to include state data in the observation + + Note that the returned observations will have a "rgbd" or "rgb" or "depth" key depending on the rgb/depth bool flags. + """ + + def __init__(self, env) -> None: + self.base_env = env.unwrapped + super().__init__(env) + new_obs = self.observation(self.base_env._init_raw_obs) + self.base_env.update_obs_space(new_obs) + + def observation(self, observation: dict): + # ret = dict() + if not hasattr(self.env, "_has_reset") or not self.env._has_reset: + self.env.reset() + self.env._has_reset = True + # observation["sensor_data"]["human_camera"] = dict(rgb=self.env.render()) + observation["sensor_data"]["base_camera"] = dict(rgb=self.env.render()) + return observation + + +class ManiSkill(EvalEnv): + INSTRUCTIONS: ClassVar[dict[str, str]] = { + "LiftPegUpright-v1": "lift the peg upright", + "PegInsertionSide-v1": "insert the peg from the side", + "PickCube-v1": "pick up the cube", + "PlugCharger-v1": "plug the charger in", + "PullCube-v1": "pull the cube towards the robot base", + "PullCubeTool-v1": "pull the cube by using the red tool", + "PushCube-v1": "push the cube away from the robot base", + "PushT-v1": "align the T shape", + "RollBall-v1": "push the ball", + "StackCube-v1": "stack the red cube on the green cube", + "PokeCube-v1": "push the cube by using the blue tool", + } + + def __init__(self, env_id, **env_kwargs): + # TODO: one could save only every nth episode by adding an episode counter which steps the record env only + # when the counter is divisible by n otherwise steps the normal env + logging.info(f"Creating ManiSkill env {env_id}") + output_dir = env_kwargs.pop("video_dir", None) + super().__init__(env_id, **env_kwargs) + logging.info(f"Created ManiSkill env {env_id}") + if "human_render_camera_configs" in env_kwargs: + self.env = HumanCameraWrapper(self.env) + + if output_dir is not None: + logging.info(f"Recording to {output_dir}") + from mani_skill.utils import wrappers + + self.env = wrappers.RecordEpisode( + self.env, + output_dir, + save_on_reset=True, + save_trajectory=True, + trajectory_name=f"eval-{env_id}", + save_video=True, + video_fps=30, + record_reward=True, + ) + logging.info(f"Done Created ManiSkill env {env_id}") + + def translate_obs(self, obs: dict[str, Any]) -> Obs: + # does not include history + return Obs( + obs={ + "default": SingleObs(cameras={"rgb_side": obs["sensor_data"]["base_camera"]["rgb"].squeeze(0).numpy()}) + }, + language_instruction=self.language_instruction, + ) + + def step(self, action: dict[str, SingleAct]) -> tuple[Obs, float, bool, bool, dict]: + # includes horizon + # careful with gripper action: the model needs to be trained on [-1, 1] interval + assert len(action) == 1, "ManiSkill expects a single robot action" + _, robot_action = next(iter(action.items())) + a = np.asarray(copy.copy(robot_action.action), dtype=np.float32) + if robot_action.gripper is not None: + a = np.concatenate([a, np.asarray([robot_action.gripper], dtype=np.float32)]) + if self.env_id == "PushT-v1": + if robot_action.gripper is not None: + a = a[:-1] + else: + a[-1] = a[-1] * 2 - 1.0 + obs, reward, success, truncated, info = self.env.step(a) + return self.translate_obs(obs), reward, success, truncated, info + + def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Obs, dict[str, Any]]: + # maniskill has a bug that does not allow None in options + obs, info = self.env.reset(seed=seed) + return self.translate_obs(obs), info + + @property + def language_instruction(self) -> str: + return self.INSTRUCTIONS[self.env_id] + + @staticmethod + def do_import(): + import mani_skill.envs + + +register_env("LiftPegUpright-v1", ManiSkill) +register_env("PegInsertionSide-v1", ManiSkill) +register_env("PickCube-v1", ManiSkill) +register_env("PlugCharger-v1", ManiSkill) +register_env("PullCube-v1", ManiSkill) +register_env("PullCubeTool-v1", ManiSkill) +register_env("PushCube-v1", ManiSkill) +register_env("PushT-v1", ManiSkill) +register_env("RollBall-v1", ManiSkill) +register_env("StackCube-v1", ManiSkill) +register_env("PokeCube-v1", ManiSkill) diff --git a/src/vlagents/eval.py b/src/vlagents/eval.py new file mode 100644 index 0000000..93fb82f --- /dev/null +++ b/src/vlagents/eval.py @@ -0,0 +1,364 @@ +import datetime +import json +import logging +import os +import shlex +import subprocess +import sys +from contextlib import contextmanager +from dataclasses import asdict +from pathlib import Path +from time import sleep +from typing import Any + +import numpy as np +from simple_slurm import Slurm +from tqdm import tqdm + +from vlagents.client import RemoteAgent +from vlagents.envs.interface import AgentConfig, EvalConfig, EvalEnv +from vlagents.policies.interface import Agent + +logging.basicConfig( + format="%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + level=logging.INFO, +) + + +def _write_camera_mp4(frames: list[np.ndarray], output_path: Path, fps: int = 30) -> None: + if not frames: + return + + height, width = frames[0].shape[:2] + process = subprocess.Popen( + [ + "ffmpeg", + "-y", + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "-s", + f"{width}x{height}", + "-r", + str(fps), + "-i", + "-", + "-an", + "-vf", + "pad=ceil(iw/2)*2:ceil(ih/2)*2", + "-c:v", + "libx264", + "-pix_fmt", + "yuv420p", + "-movflags", + "+faststart", + str(output_path), + ], + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + assert process.stdin is not None + for frame in frames: + process.stdin.write(np.ascontiguousarray(frame).astype(np.uint8).tobytes()) + process.stdin.close() + process.wait() + + +def single_eval( + env: EvalEnv, agent: Agent, max_steps: int, ith_episode: int, start_seed: int +) -> tuple[list[float], list[float], list[float]]: + logging.debug("Starting evaluation") + obs, _ = env.reset(seed=start_seed + ith_episode) # ensure different seed for each episode + if obs.language_instruction is None: + obs.language_instruction = env.language_instruction + single_obs = next(iter(obs.obs.values())) + cameras = single_obs.info.pop("high_res_cameras", single_obs.cameras) + logging.debug("Reset env") + done = False + truncated = False + step = 0.0 + rewards = [] + im = [] + while not done and not truncated and max_steps > step: + if obs.language_instruction is None: + obs.language_instruction = env.language_instruction + obs, reward, done, truncated, _ = env.chunk_step(agent.act(obs), max_steps=max_steps - int(step)) + if obs.language_instruction is None: + obs.language_instruction = env.language_instruction + single_obs = next(iter(obs.obs.values())) + cameras = single_obs.info.pop("high_res_cameras", single_obs.cameras) + reward = float(reward) + done, truncated = bool(done), bool(truncated) + step += env.last_chunk_steps + rewards.append(reward) + im.append(cameras) + + cam_path = os.environ.get("CAM_PATH", None) + if cam_path is not None and im: + output_dir = Path(os.environ["CAM_PATH"]) / env.env_id + output_dir.mkdir(exist_ok=True, parents=True) + timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + for camera in im[0].keys(): + _write_camera_mp4( + [img[camera] for img in im], + output_dir / f"{ith_episode}_{camera}_{timestamp}.mp4", + ) + + env.reset() + logging.debug(f"Finished evaluation with {step} steps and reward {reward}, success {done}") + # success, last reward and number of steps + return done, rewards, step + + +per_process_cache = {} + + +def create_env_agent(agent_config: AgentConfig, cfg: EvalConfig) -> tuple[EvalEnv, RemoteAgent]: + logging.debug(f"retrieving env {cfg.env_id} and agent") + key = (cfg.env_id, agent_config.host, agent_config.port) + if key not in per_process_cache: + logging.info(f"env {cfg.env_id} not available, creating new env and agent") + env = EvalEnv.from_id(cfg.env_id, execution_horizon=cfg.execution_horizon, **cfg.env_kwargs) + logging.info("done creating env") + agent = RemoteAgent( + agent_config.host, + agent_config.port, + agent_config.agent_name, + on_same_machine=cfg.same_machine, + jpeg_encoding=cfg.jpeg_encoding, + image_size=cfg.image_size, + ) + logging.info("done creating agent") + per_process_cache[key] = (env, agent) + return per_process_cache[key] + + +def run_episode(args: tuple[int, list[EvalConfig], int, AgentConfig]) -> tuple[list[float], list[float], list[float]]: + i, cfgs, episodes, agent_cfg = args + cfg = cfgs[i // episodes] + env, agent = create_env_agent(agent_cfg, cfg) + # busy wait for server to finish initialization + while not agent.is_initialized(): + logging.info("Waiting for agent to initialize...") + sleep(5) + return single_eval(env, agent, cfg.max_steps_per_episode, i, start_seed=cfg.seed) + + +def multi_eval( + agent_cfg: AgentConfig, cfgs: list[EvalConfig], episodes: int = 100 +) -> tuple[np.ndarray, list[list[list[float]]]]: + # return is [envs, episodes, 3(success, reward, steps)], [envs, episodes, rewards for all steps in the episode] + logging.info(f"Starting evaluation with {len(cfgs)} environments and {episodes} episodes each") + + # np.random.seed(cfgs[0].seed) + args = [(i, cfgs, episodes, agent_cfg) for i in range(len(cfgs) * episodes)] + single_results = [run_episode(arg) for arg in tqdm(args)] + + single_results_last_reward = np.array([(i[0], i[1][-1], i[2]) for i in single_results]) + + # this works because row-major order + # per_env_results = single_results.reshape(len(cfgs), episodes, 3) + per_env_results_last_reward = single_results_last_reward.reshape(len(cfgs), episodes, 3) + per_env_results_rewards = [ + [i[1] for i in single_results[i : i + episodes]] for i in range(0, len(single_results), episodes) + ] + return per_env_results_last_reward, per_env_results_rewards + + +@contextmanager +def start_server( + agent_name: str, + kwargs: dict[str, Any], + port: int = 8080, + host: str = "localhost", + python_path: str = sys.executable, +): + """Start the agent server in a subprocess as a context manager. + + This ensures that the server is properly stopped when exiting the context and + that all logs are printed to the console. + + Args: + agent_name (str): Name of the agent to start. + kwargs (dict[str, Any]): Additional keyword arguments for the agent. + port (int): Port to start the server on. Defaults to 8080. + host (str): Host to bind the server to. Defaults to "localhost". + python_path (str): Path to the Python interpreter to use. If you use conda you can look up the path with `conda info --envs`. + It can also be a format string that will be formatted with the agent_name, e.g. "conda run -n {agent_name} python". + Defaults to the current interpreter (`sys.executable`). + """ + cmd = [ + python_path.format(agent_name=agent_name), + "-m", + "vlagents", + "start-server", + f"{agent_name}", + f"--port={port}", + f"--host={host}", + f"--kwargs={json.dumps(kwargs)}", + ] + logging.info("Server starting: %s", " ".join(cmd)) + env = os.environ.copy() + source_root = str(Path(__file__).resolve().parents[1]) + existing_pythonpath = env.get("PYTHONPATH") + env["PYTHONPATH"] = source_root if not existing_pythonpath else f"{source_root}:{existing_pythonpath}" + p = subprocess.Popen(cmd, env=env) + try: + yield p + finally: + # Stop the server no matter how we exit the with-block (success or exception). + try: + p.send_signal(subprocess.signal.SIGINT) + p.wait(timeout=5) + except Exception: + pass + if p.poll() is None: + p.terminate() + try: + p.wait(timeout=3) + except subprocess.TimeoutExpired: + p.kill() + logging.info("Server stopped") + + +def evaluation( + agent_cfg: AgentConfig, + eval_cfgs: list[EvalConfig], + episodes: int = 100, +): + per_process_cache.clear() + logging.info(f"Starting evaluation with {agent_cfg.agent_name} and {agent_cfg.agent_kwargs}") + try: + with start_server( + agent_cfg.agent_name, agent_cfg.agent_kwargs, agent_cfg.port, agent_cfg.host, agent_cfg.python_path + ): + sleep(30) + res = multi_eval(agent_cfg, eval_cfgs, episodes) + except Exception: + # Ensures you SEE the client's stack trace and any logged errors. + logging.exception("Client failed") + raise + + logging.info(f"Results (success, reward, steps) for all envs: {res[0].mean(axis=1)}") + logging.info( + f"Mean reward for all envs: {[np.mean([np.mean(ep_rewards) for ep_rewards in env_rewards]) for env_rewards in res[1]]}" + ) + # print indices of successful episodes + for idx, env in enumerate(res[0]): + logging.info(f"Env {eval_cfgs[idx].env_id} successful episodes: {np.where(env[:, 0])[0]}") + return res + + +def run_eval( + agent_cfg: AgentConfig, + eval_cfgs: list[EvalConfig], + wandb_entity: str, + wandb_project: str, + wandb_note: str, + wandb_name: str, + checkpoint_steps: list[int], + slurm: Slurm, + output_path: str, + wandb_group: str | None = None, + episodes: int = 100, + n_processes: int | None = None, + n_gpus: int = 1, + python_path: str = "python", +): + eval_cmd = shlex.quote( + shlex.join( + [ + "-m", + "vlagents", + "run-eval", + f"--agent-cfg={json.dumps(asdict(agent_cfg))}", + f"--episodes={episodes}", + f"--n-processes={n_processes}", + f"--eval-cfgs={json.dumps([asdict(cfg) for cfg in eval_cfgs])}", + f"--wandb-group={wandb_group.replace(':', '_') if wandb_group else ''}", + f"--wandb-project={wandb_project}", + f"--wandb-entity={wandb_entity}", + f"--wandb-note={wandb_note}", + f"--wandb-name={wandb_name}", + f"--n-gpus={n_gpus}", + f"--steps={json.dumps(checkpoint_steps)}", + f"--output-path={output_path}", + ] + ) + ) + + python_path += eval_cmd + slurm.sbatch(python_path) + + +def write_results( + results: np.ndarray, + rewards: list[list[list[float]]], + eval_cfgs: list[EvalConfig], + agent_cfg: AgentConfig, + out: str = "", + grouped_eval_cfgs: list[list[EvalConfig]] | None = None, +) -> str: + # first read json, if not exists write empty list + path = os.path.join(out, f"results_{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.json") + if not os.path.exists(path): + with open(path, "w") as f: + json.dump([], f) + with open(path, "r") as f: + prev_results = json.load(f) + assert isinstance(prev_results, list) + + flatten_rewards = [[item for sublist in env_rewards for item in sublist] for env_rewards in rewards] + mean_rewards = [np.mean(env_rewards) for env_rewards in flatten_rewards] + grouped_eval_cfgs = grouped_eval_cfgs or [[cfg] for cfg in eval_cfgs] + + for idx, (cfg, cfg_group) in enumerate(zip(eval_cfgs, grouped_eval_cfgs, strict=True)): + success_mean, reward_mean, steps_mean = results[idx].mean(axis=0, keepdims=False) + success_max, reward_max, steps_max = results[idx].max(axis=0, keepdims=False) + success_min, reward_min, steps_min = results[idx].min(axis=0, keepdims=False) + sucess_std, reward_std, steps_std = results[idx].std(axis=0, keepdims=False) + success_median, reward_median, steps_median = np.median(results[idx], axis=0, keepdims=False) + result_entry = { + "success": { + "mean": success_mean, + "max": success_max, + "min": success_min, + "std": sucess_std, + "median": success_median, + "values": results[idx, :, 0].tolist(), + }, + "reward_last_step": { + "mean": reward_mean, + "max": reward_max, + "min": reward_min, + "std": reward_std, + "median": reward_median, + "values": results[idx, :, 1].tolist(), + }, + "rewards": { + "mean": mean_rewards[idx], + "values": rewards[idx], + }, + "steps": { + "mean": steps_mean, + "max": steps_max, + "min": steps_min, + "std": steps_std, + "median": steps_median, + "values": results[idx, :, 2].tolist(), + }, + "episodes": results.shape[1], + "timestamp": datetime.datetime.now().isoformat(), + "env_cfg": asdict(cfg), + "agent_cfg": asdict(agent_cfg), + } + if len(cfg_group) > 1: + result_entry["merged_env_cfgs"] = [asdict(group_cfg) for group_cfg in cfg_group] + prev_results.append(result_entry) + + with open(path, "w") as f: + json.dump(prev_results, f, indent=2) + return path diff --git a/src/vlagents/evaluator_envs.py b/src/vlagents/evaluator_envs.py deleted file mode 100644 index 9b41630..0000000 --- a/src/vlagents/evaluator_envs.py +++ /dev/null @@ -1,686 +0,0 @@ -import copy -import datetime -import json -import logging -import os -import shlex -import subprocess -from abc import ABC -from contextlib import contextmanager -from dataclasses import asdict, dataclass -from multiprocessing import Pool -from pathlib import Path -from time import sleep -from typing import Any - -import gymnasium as gym -import numpy as np -from PIL import Image -from simple_slurm import Slurm -from tqdm import tqdm - -from vlagents.client import RemoteAgent -from vlagents.policies import Act, Agent, Obs -from vlagents.wrappers import HumanCameraWrapper - -logging.basicConfig( - format="%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - level=logging.INFO, -) - - -class EvaluatorEnv(ABC): - ENVS: dict[str, "EvaluatorEnv"] = {} - - def __init__(self, env_id: str, **env_kwargs) -> None: - self.do_import() - self.env = gym.make(env_id, **env_kwargs) - self.env_id = env_id - - def step(self, action: Act) -> tuple[Obs, float, bool, bool, dict]: - raise NotImplementedError - - def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Obs, dict[str, Any]]: - raise NotImplementedError - - @property - def language_instruction(self) -> str: - raise NotImplementedError - - @staticmethod - def register(env_id: str, env: "EvaluatorEnv") -> None: - EvaluatorEnv.ENVS[env_id] = env - - @staticmethod - def make(env_id: str, **env_kwargs) -> "EvaluatorEnv": - return EvaluatorEnv.ENVS[env_id](env_id, **env_kwargs) - - @staticmethod - def do_import(): - raise NotImplementedError - - -class RCSDuoBench(EvaluatorEnv): - INSTRUCTIONS = {} - - def __init__(self, env_id, **env_kwargs): - self.robot_keys: str = env_kwargs.pop("robot_keys", ["left", "right"]) - self.control_mode: str = env_kwargs.pop("control_mode", "joints") - self._instruction: str | None = None - super().__init__(env_id, **env_kwargs) - - def translate_obs(self, obs: dict[str, Any]) -> Obs: - cameras = {} - for key in obs["frames"]: - cameras[key] = obs["frames"][key]["rgb"]["data"] - cameras[key] = np.array(Image.fromarray(cameras[key]).resize((224, 224), Image.Resampling.BILINEAR)) - state = [] - for key in self.robot_keys: - state.append(obs[key]["joints"]) - state.append(obs[key]["gripper"]) - - return Obs( - cameras=cameras, - gripper=None, - state=np.concatenate(state), - info={"high_res_cameras": {key: obs["frames"][key]["rgb"]["data"] for key in obs["frames"]}}, - ) - - def step(self, action: Act) -> tuple[Obs, float, bool, bool, dict]: - assert ( - len(action.action.shape) == 1 - ), "this function cannot deal with batches or action chunks, please return single actions" - env_action = {} - for idx, robot in enumerate(self.robot_keys): - if self.control_mode == "joints": - env_action[robot] = { - "joints": action.action[idx * 8 : idx * 8 + 7], - "gripper": action.action[idx * 8 + 7 : idx * 8 + 8], - } - else: - env_action[robot] = { - "xyzrpy": action.action[idx * 7 : idx * 7 + 6], - "gripper": action.action[idx * 7 + 6 : idx * 7 + 7], - } - obs, reward, success, truncated, info = self.env.step(env_action) - r = float(reward) - - return self.translate_obs(obs), r, success, truncated, info - - def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Obs, dict[str, Any]]: - obs, info = self.env.reset(seed=seed, options=options) - self._instruction = info["instruction"] - return self.translate_obs(obs), info - - @property - def language_instruction(self) -> str: - assert self._instruction is not None - return self._instruction - - @staticmethod - def do_import(): - import rcs - from rcs_duobench.tasks import ( - ball_maze, - bin_sort, - block_balance, - carry_pot, - hinge_chest, - join_blocks, - pour_marbles, - spring_door, - transfer_cube, - transfer_gate, - transfer_reorient, - ) - - -EvaluatorEnv.register("duobench/ball_maze", RCSDuoBench) -EvaluatorEnv.register("duobench/bin_sort", RCSDuoBench) -EvaluatorEnv.register("duobench/block_balance", RCSDuoBench) -EvaluatorEnv.register("duobench/carry_pot", RCSDuoBench) -EvaluatorEnv.register("duobench/join_blocks", RCSDuoBench) -EvaluatorEnv.register("duobench/hinge_chest", RCSDuoBench) -EvaluatorEnv.register("duobench/pour_marbles", RCSDuoBench) -EvaluatorEnv.register("duobench/spring_door", RCSDuoBench) -EvaluatorEnv.register("duobench/transfer_cube", RCSDuoBench) -EvaluatorEnv.register("duobench/transfer_gate", RCSDuoBench) -EvaluatorEnv.register("duobench/transfer_reorient", RCSDuoBench) - - -class ManiSkill(EvaluatorEnv): - INSTRUCTIONS = { - "LiftPegUpright-v1": "lift the peg upright", - "PegInsertionSide-v1": "insert the peg from the side", - "PickCube-v1": "pick up the cube", - "PlugCharger-v1": "plug the charger in", - "PullCube-v1": "pull the cube towards the robot base", - "PullCubeTool-v1": "pull the cube by using the red tool", - "PushCube-v1": "push the cube away from the robot base", - "PushT-v1": "align the T shape", - "RollBall-v1": "push the ball", - "StackCube-v1": "stack the red cube on the green cube", - "PokeCube-v1": "push the cube by using the blue tool", - } - - def __init__(self, env_id, **env_kwargs): - # TODO: one could save only every nth episode by adding an episode counter which steps the record env only - # when the counter is divisible by n otherwise steps the normal env - logging.info(f"Creating ManiSkill env {env_id}") - output_dir = env_kwargs.pop("video_dir", None) - super().__init__(env_id, **env_kwargs) - logging.info(f"Created ManiSkill env {env_id}") - if "human_render_camera_configs" in env_kwargs: - self.env = HumanCameraWrapper(self.env) - - if output_dir is not None: - logging.info(f"Recording to {output_dir}") - from mani_skill.utils import wrappers - - self.env = wrappers.RecordEpisode( - self.env, - output_dir, - save_on_reset=True, - save_trajectory=True, - trajectory_name=f"eval-{env_id}", - save_video=True, - video_fps=30, - record_reward=True, - ) - logging.info(f"Done Created ManiSkill env {env_id}") - - def translate_obs(self, obs: dict[str, Any]) -> Obs: - # does not include history - return Obs( - cameras=dict(rgb_side=obs["sensor_data"]["base_camera"]["rgb"].squeeze(0).numpy()), - # gripper=float(not obs["extra"]["is_grasped"]), - ) - - def step(self, action: Act) -> tuple[Obs, float, bool, bool, dict]: - # includes horizon - # careful with gripper action: the model needs to be trained on [-1, 1] interval - - a = copy.copy(action.action[0]) - # a[-1] = -1.0 if a[-1] < 0.9 else 1.0 - if self.env_id == "PushT-v1": - a = a[:-1] - else: - a[-1] = a[-1] * 2 - 1.0 - obs, reward, success, truncated, info = self.env.step(a) - return self.translate_obs(obs), reward, success, truncated, info - - def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Obs, dict[str, Any]]: - # maniskill has a bug that does not allow None in options - obs, info = self.env.reset(seed=seed) - return self.translate_obs(obs), info - - @property - def language_instruction(self) -> str: - return self.INSTRUCTIONS[self.env_id] - - @staticmethod - def do_import(): - import mani_skill.envs - - -EvaluatorEnv.register("LiftPegUpright-v1", ManiSkill) -EvaluatorEnv.register("PegInsertionSide-v1", ManiSkill) -EvaluatorEnv.register("PickCube-v1", ManiSkill) -EvaluatorEnv.register("PlugCharger-v1", ManiSkill) -EvaluatorEnv.register("PullCube-v1", ManiSkill) -EvaluatorEnv.register("PullCubeTool-v1", ManiSkill) -EvaluatorEnv.register("PushCube-v1", ManiSkill) -EvaluatorEnv.register("PushT-v1", ManiSkill) -EvaluatorEnv.register("RollBall-v1", ManiSkill) -EvaluatorEnv.register("StackCube-v1", ManiSkill) -EvaluatorEnv.register("PokeCube-v1", ManiSkill) - - -class Libero(EvaluatorEnv): - - def __init__(self, env_id: str, reset_steps: int = 14, **env_kwargs) -> None: - """ - For supported env_kwargs checkout ControlEnv class in libero. - We add the following env_kwargs on top: - - task_id (int): libero task id for given task suite. The number of tasks per task suite can checked with Libero.n_tasks(env_id). Defaults to 0. - - control_mode (str): either 'relative' or 'absolute'. Defaults to 'relative'. - - """ - logging.info("Creating Libero env") - self.env_kwargs = env_kwargs - self.reset_steps = reset_steps - self.control_mode = self.env_kwargs.pop("control_mode", "relative") - self.env, self._language_instruction, self.task_name, self.task_suite, self.task_id, self.task = self._make_gym( - env_id, **self.env_kwargs - ) - logging.info( - f"Created Libero env, task suite: {env_id}, task id: {self.task_id}, task name {self.task_name}, instruction: {self._language_instruction}" - ) - self.env_id = env_id - - @staticmethod - def n_tasks(env_id: str) -> int: - from libero.libero import benchmark, get_libero_path - - benchmark_dict = benchmark.get_benchmark_dict() - task_suite = benchmark_dict[env_id]() - return task_suite.n_tasks - - @staticmethod - def _make_gym(env_id, **env_kwargs): - from libero.libero import benchmark, get_libero_path - from libero.libero.envs import OffScreenRenderEnv - - benchmark_dict = benchmark.get_benchmark_dict() - - task_suite = benchmark_dict[env_id]() - task_id = min(max(env_kwargs.pop("task_id", 0), 0), task_suite.n_tasks - 1) - task = task_suite.get_task(task_id) - - task_bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file) - env = OffScreenRenderEnv( - bddl_file_name=task_bddl_file, - **env_kwargs, - ) - - return env, task.language, task.name, task_suite, task_id, task - - def translate_obs(self, obs: dict[str, Any]) -> Obs: - return Obs( - cameras=dict(rgb_side=obs["agentview_image"][::-1], rgb_wrist=obs["robot0_eye_in_hand_image"][::-1]), - gripper=obs["robot0_gripper_qpos"] / 0.04, # normalize - ) - - def step(self, action: Act) -> tuple[Obs, float, bool, bool, dict]: - # change gripper to libero format (-1, 1) where -1 is open - act = np.copy(action.action) - act[-1] = (1 - act[-1]) * 2 - 1.0 - obs, reward, done, info = self.env.step(act) - success = self.env.check_success() - return self.translate_obs(obs), reward, success, done, info - - def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Obs, dict[str, Any]]: - if seed is not None: - self.env.seed(seed) - obs = self.env.reset() - init_states = self.task_suite.get_task_init_states( - self.task_id - ) # for benchmarking purpose, we fix the a set of initial states - init_state_id = 0 - self.env.set_init_state(init_states[init_state_id]) - - for robot in self.env.robots: - robot.controller.use_delta = True - for _ in range(self.reset_steps): - # steps the environment to filter out falling objects - obs, _, _, _ = self.env.step( - np.zeros(8) if "JOINT" in self.env_kwargs.get("controller", "OSC_POSE") else np.zeros(7) - ) - - if self.control_mode == "absolute": - for robot in self.env.robots: - robot.controller.use_delta = False - elif self.control_mode == "relative": - for robot in self.env.robots: - robot.controller.use_delta = True - else: - raise ValueError(f"Invalid control mode: {self.control_mode}, use 'absolute' or 'relative'.") - - return self.translate_obs(obs), {} - - @property - def language_instruction(self) -> str: - return self._language_instruction - - -EvaluatorEnv.register("libero_10", Libero) -EvaluatorEnv.register("libero_90", Libero) -EvaluatorEnv.register("libero_100", Libero) -EvaluatorEnv.register("libero_spatial", Libero) -EvaluatorEnv.register("libero_object", Libero) -EvaluatorEnv.register("libero_goal", Libero) - - -@dataclass -class EvalConfig: - env_id: str - env_kwargs: dict[str, Any] - max_steps_per_episode: int = 100 - seed: int = 42 - same_machine: bool = False - jpeg_encoding: bool = False - - -@dataclass -class AgentConfig: - host: str - agent_name: str - agent_kwargs: dict[str, Any] - python_path: str = "python" - """modify this if you want to use a specific python environment """ - port: int = 8080 - - -def _write_camera_mp4(frames: list[np.ndarray], output_path: Path, fps: int = 30) -> None: - if not frames: - return - - height, width = frames[0].shape[:2] - process = subprocess.Popen( - [ - "ffmpeg", - "-y", - "-f", - "rawvideo", - "-pix_fmt", - "rgb24", - "-s", - f"{width}x{height}", - "-r", - str(fps), - "-i", - "-", - "-an", - "-vf", - "pad=ceil(iw/2)*2:ceil(ih/2)*2", - "-c:v", - "libx264", - "-pix_fmt", - "yuv420p", - "-movflags", - "+faststart", - str(output_path), - ], - stdin=subprocess.PIPE, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - assert process.stdin is not None - for frame in frames: - process.stdin.write(np.ascontiguousarray(frame).astype(np.uint8).tobytes()) - process.stdin.close() - process.wait() - - -def single_eval( - env: EvaluatorEnv, agent: Agent, max_steps: int, ith_episode: int, start_seed: int -) -> tuple[list[float], list[float], list[float]]: - logging.debug(f"Starting evaluation") - obs, _ = env.reset(seed=start_seed + ith_episode) # ensure different seed for each episode - cameras = obs.info.pop("high_res_cameras", obs.cameras) - logging.debug(f"Reset env") - agent.reset(copy.deepcopy(obs), env.language_instruction) - logging.debug(f"Reset agent") - done = False - truncated = False - step = 0.0 - rewards = [] - im = [] - while not done and not truncated and max_steps > step: - action = agent.act(obs) - obs, reward, done, truncated, _ = env.step(action) - cameras = obs.info.pop("high_res_cameras", obs.cameras) - reward = float(reward) - done, truncated = bool(done), bool(truncated) - step += 1 - rewards.append(reward) - im.append(cameras) - - cam_path = os.environ.get("CAM_PATH", None) - if cam_path is not None and im: - output_dir = Path(os.environ["CAM_PATH"]) / env.env_id - output_dir.mkdir(exist_ok=True, parents=True) - timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - for camera in im[0].keys(): - _write_camera_mp4( - [img[camera] for img in im], - output_dir / f"{ith_episode}_{camera}_{timestamp}.mp4", - ) - - env.reset() - logging.debug(f"Finished evaluation with {step} steps and reward {reward}, success {done}") - # success, last reward and number of steps - return done, rewards, step - - -per_process_cache = {} - - -def create_env_agent(agent_config: AgentConfig, cfg: EvalConfig) -> tuple[EvaluatorEnv, RemoteAgent]: - logging.debug(f"retrieving env {cfg.env_id} and agent") - key = (cfg.env_id, agent_config.host, agent_config.port) - if key not in per_process_cache: - logging.info(f"env {cfg.env_id} not available, creating new env and agent") - env = EvaluatorEnv.make(cfg.env_id, **cfg.env_kwargs) - logging.info("done creating env") - agent = RemoteAgent( - agent_config.host, - agent_config.port, - agent_config.agent_name, - on_same_machine=cfg.same_machine, - jpeg_encoding=cfg.jpeg_encoding, - ) - logging.info("done creating agent") - per_process_cache[key] = (env, agent) - return per_process_cache[key] - - -def run_episode(args: tuple[int, list[EvalConfig], int, AgentConfig]) -> tuple[list[float], list[float], list[float]]: - i, cfgs, episodes, agent_cfg = args - cfg = cfgs[i // episodes] - env, agent = create_env_agent(agent_cfg, cfg) - # busy wait for server to finish initialization - while not agent.is_initialized(): - logging.info("Waiting for agent to initialize...") - sleep(5) - return single_eval(env, agent, cfg.max_steps_per_episode, i, start_seed=cfg.seed) - - -def multi_eval( - agent_cfg: AgentConfig, cfgs: list[EvalConfig], episodes: int = 100 -) -> tuple[np.ndarray, list[list[list[float]]]]: - # return is [envs, episodes, 3(success, reward, steps)], [envs, episodes, rewards for all steps in the episode] - logging.info(f"Starting evaluation with {len(cfgs)} environments and {episodes} episodes each") - - # np.random.seed(cfgs[0].seed) - args = [(i, cfgs, episodes, agent_cfg) for i in range(len(cfgs) * episodes)] - single_results = [run_episode(arg) for arg in tqdm(args)] - - single_results_last_reward = np.array([(i[0], i[1][-1], i[2]) for i in single_results]) - - # this works because row-major order - # per_env_results = single_results.reshape(len(cfgs), episodes, 3) - per_env_results_last_reward = single_results_last_reward.reshape(len(cfgs), episodes, 3) - per_env_results_rewards = [ - [i[1] for i in single_results[i : i + episodes]] for i in range(0, len(single_results), episodes) - ] - return per_env_results_last_reward, per_env_results_rewards - - -@contextmanager -def start_server( - agent_name: str, kwargs: dict[str, Any], port: int = 8080, host: str = "localhost", python_path: str = "python" -): - """Start the agent server in a subprocess as a context manager. - - This ensures that the server is properly stopped when exiting the context and - that all logs are printed to the console. - - Args: - agent_name (str): Name of the agent to start. - kwargs (dict[str, Any]): Additional keyword arguments for the agent. - port (int): Port to start the server on. Defaults to 8080. - host (str): Host to bind the server to. Defaults to "localhost". - python_path (str): Path to the Python interpreter to use. If you use conda you can look up the path with `conda info --envs`. - It can also be a format string that will be formatted with the agent_name, e.g. "conda run -n {agent_name} python". - Defaults to "python". - """ - cmd = [ - python_path.format(agent_name=agent_name), - "-m", - "vlagents", - "start-server", - f"{agent_name}", - f"--port={port}", - f"--host={host}", - f"--kwargs={json.dumps(kwargs)}", - ] - logging.info("Server starting: %s", " ".join(cmd)) - p = subprocess.Popen(cmd) - try: - yield p - finally: - # Stop the server no matter how we exit the with-block (success or exception). - try: - p.send_signal(subprocess.signal.SIGINT) - p.wait(timeout=5) - except Exception: - pass - if p.poll() is None: - p.terminate() - try: - p.wait(timeout=3) - except subprocess.TimeoutExpired: - p.kill() - logging.info("Server stopped") - - -def evaluation( - agent_cfg: AgentConfig, - eval_cfgs: list[EvalConfig], - episodes: int = 100, -): - per_process_cache.clear() - logging.info(f"Starting evaluation with {agent_cfg.agent_name} and {agent_cfg.agent_kwargs}") - try: - with start_server( - agent_cfg.agent_name, agent_cfg.agent_kwargs, agent_cfg.port, agent_cfg.host, agent_cfg.python_path - ): - sleep(30) - res = multi_eval(agent_cfg, eval_cfgs, episodes) - except Exception: - # Ensures you SEE the client's stack trace and any logged errors. - logging.exception("Client failed") - raise - - logging.info(f"Results (success, reward, steps) for all envs: {res[0].mean(axis=1)}") - logging.info( - f"Mean reward for all envs: {[np.mean([np.mean(ep_rewards) for ep_rewards in env_rewards]) for env_rewards in res[1]]}" - ) - # print indices of successful episodes - for idx, env in enumerate(res[0]): - logging.info(f"Env {eval_cfgs[idx].env_id} successful episodes: {np.where(env[:, 0])[0]}") - return res - - -def run_eval( - agent_cfg: AgentConfig, - eval_cfgs: list[EvalConfig], - wandb_entity: str, - wandb_project: str, - wandb_note: str, - wandb_name: str, - checkpoint_steps: list[int], - slurm: Slurm, - output_path: str, - wandb_group: str | None = None, - episodes: int = 100, - n_processes: int | None = None, - n_gpus: int = 1, - python_path: str = "python", -): - eval_cmd = shlex.quote( - shlex.join( - [ - "-m", - "vlagents", - "run-eval", - f"--agent-cfg={json.dumps(asdict(agent_cfg))}", - f"--episodes={episodes}", - f"--n-processes={n_processes}", - f"--eval-cfgs={json.dumps([asdict(cfg) for cfg in eval_cfgs])}", - f"--wandb-group={wandb_group.replace(':', '_') if wandb_group else ''}", - f"--wandb-project={wandb_project}", - f"--wandb-entity={wandb_entity}", - f"--wandb-note={wandb_note}", - f"--wandb-name={wandb_name}", - f"--n-gpus={n_gpus}", - f"--steps={json.dumps(checkpoint_steps)}", - f"--output-path={output_path}", - ] - ) - ) - - python_path += eval_cmd - slurm.sbatch(python_path) - - -def write_results( - results: np.ndarray, - rewards: list[list[list[float]]], - eval_cfgs: list[EvalConfig], - agent_cfg: AgentConfig, - out: str = "", - grouped_eval_cfgs: list[list[EvalConfig]] | None = None, -) -> str: - # first read json, if not exists write empty list - path = os.path.join(out, f"results_{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.json") - if not os.path.exists(path): - with open(path, "w") as f: - json.dump([], f) - with open(path, "r") as f: - prev_results = json.load(f) - assert isinstance(prev_results, list) - - flatten_rewards = [[item for sublist in env_rewards for item in sublist] for env_rewards in rewards] - mean_rewards = [np.mean(env_rewards) for env_rewards in flatten_rewards] - grouped_eval_cfgs = grouped_eval_cfgs or [[cfg] for cfg in eval_cfgs] - - for idx, (cfg, cfg_group) in enumerate(zip(eval_cfgs, grouped_eval_cfgs, strict=True)): - success_mean, reward_mean, steps_mean = results[idx].mean(axis=0, keepdims=False) - success_max, reward_max, steps_max = results[idx].max(axis=0, keepdims=False) - success_min, reward_min, steps_min = results[idx].min(axis=0, keepdims=False) - sucess_std, reward_std, steps_std = results[idx].std(axis=0, keepdims=False) - success_median, reward_median, steps_median = np.median(results[idx], axis=0, keepdims=False) - result_entry = { - "success": { - "mean": success_mean, - "max": success_max, - "min": success_min, - "std": sucess_std, - "median": success_median, - "values": results[idx, :, 0].tolist(), - }, - "reward_last_step": { - "mean": reward_mean, - "max": reward_max, - "min": reward_min, - "std": reward_std, - "median": reward_median, - "values": results[idx, :, 1].tolist(), - }, - "rewards": { - "mean": mean_rewards[idx], - "values": rewards[idx], - }, - "steps": { - "mean": steps_mean, - "max": steps_max, - "min": steps_min, - "std": steps_std, - "median": steps_median, - "values": results[idx, :, 2].tolist(), - }, - "episodes": results.shape[1], - "timestamp": datetime.datetime.now().isoformat(), - "env_cfg": asdict(cfg), - "agent_cfg": asdict(agent_cfg), - } - if len(cfg_group) > 1: - result_entry["merged_env_cfgs"] = [asdict(group_cfg) for group_cfg in cfg_group] - prev_results.append(result_entry) - - with open(path, "w") as f: - json.dump(prev_results, f, indent=2) - return path diff --git a/src/vlagents/policies.py b/src/vlagents/policies.py deleted file mode 100644 index 932f31e..0000000 --- a/src/vlagents/policies.py +++ /dev/null @@ -1,826 +0,0 @@ -import base64 -import copy -import json -import logging -import os -from collections import deque -from dataclasses import dataclass, field -from functools import partial, reduce -from multiprocessing import resource_tracker, shared_memory -from operator import getitem -from pathlib import Path -from typing import Any, Union - -import numpy as np -import simplejpeg -from PIL import Image - - -@dataclass(kw_only=True) -class SharedMemoryPayload: - shm_name: str - shape: tuple[int, ...] - dtype: str = "uint8" - - -class CameraDataType: - SHARED_MEMORY = "shared_memory" - JPEG_ENCODED = "jpeg_encoded" - RAW = "raw" - - -@dataclass(kw_only=True) -class Obs: - cameras: dict[str, np.ndarray | SharedMemoryPayload | str] = field(default_factory=dict) - camera_data_type: str = CameraDataType.RAW - gripper: float | None = None - # TODO: add context about what the state means, and its dimensions - # theoratically it would be joints, xyzrpy and absolute or relative - state: np.ndarray | None = None - info: dict[str, Any] = field(default_factory=dict) - - -@dataclass(kw_only=True) -class Act: - action: np.ndarray - done: bool = False - info: dict[str, Any] = field(default_factory=dict) - - -class Agent: - def __init__( - self, default_checkpoint_path: str, checkpoint_path: str | None = None, checkpoint_step: int | None = None - ) -> None: - self.instruction = None - self.step = -1 - self.episode = -1 - self.checkpoint_step = checkpoint_step - self.default_checkpoint_path = default_checkpoint_path - self.checkpoint_path = checkpoint_path - self._shm: dict[str, shared_memory.SharedMemory] = {} - - def initialize(self): - # heavy initialization, e.g. loading models - pass - - def _to_numpy(self, obs: Obs) -> Obs: - """transparently uses shared memory if configured and modifies obs in place""" - if obs.camera_data_type == CameraDataType.SHARED_MEMORY: - camera_dict = {} - for camera_name, camera_data in obs.cameras.items(): - assert isinstance(camera_data, SharedMemoryPayload) - if camera_data.shm_name not in self._shm: - self._shm[camera_data.shm_name] = shared_memory.SharedMemory(camera_data.shm_name) - camera_dict[camera_name] = np.ndarray( - camera_data.shape, dtype=camera_data.dtype, buffer=self._shm[camera_data.shm_name].buf - ) - obs.cameras = camera_dict - elif obs.camera_data_type == CameraDataType.JPEG_ENCODED: - camera_dict = {} - for camera_name, camera_data in obs.cameras.items(): - assert isinstance(camera_data, str) - camera_dict[camera_name] = simplejpeg.decode_jpeg(base64.urlsafe_b64decode(camera_data)) - obs.cameras = camera_dict - obs.camera_data_type = CameraDataType.RAW - return obs - - def act(self, obs: Obs) -> Act: - assert self.instruction is not None, "forgot reset?" - self.step += 1 - self._to_numpy(obs) - - return Act(action=np.zeros(7, dtype=np.float32), done=False, info={}) - - def reset(self, obs: Obs, instruction: Any, **kwargs) -> dict[str, Any]: - logging.info(f"Resetting agent, new instruction: {instruction} ###############") - self.step = 0 - self.episode += 1 - self.instruction = instruction - self._to_numpy(obs) - # info - return {} - - def __enter__(self): - pass - - def __exit__(self, *args, **kwargs): - self.close() - - def close(self, *args, **kwargs): - for shm in self._shm.values(): - shm.close() - resource_tracker.unregister(shm._name, "shared_memory") - self._shm = {} - - -class TestAgent(Agent): - - def __init__(self, **kwargs) -> None: - super().__init__(default_checkpoint_path="", **kwargs) - self.i = 0 - - def act(self, obs: Obs) -> Act: - super().act(obs) - # echo data back for testing - info = { - "shapes": {k: v.shape for k, v in obs.cameras.items()}, - "dtype": {k: v.dtype.name for k, v in obs.cameras.items()}, - "data": {k: v for k, v in obs.cameras.items()}, - } - a = Act(action=np.array([0, 0, 0, 0, 0, 0, self.i % 2], dtype=np.float32), done=False, info=info) - self.i += 1 - return a - - def reset(self, obs: Obs, instruction: Any, **kwargs) -> dict[str, Any]: - super().reset(obs, instruction, **kwargs) - info = { - "shapes": {k: v.shape for k, v in obs.cameras.items()}, - "dtype": {k: v.dtype.name for k, v in obs.cameras.items()}, - "data": {k: v for k, v in obs.cameras.items()}, - "instruction": instruction, - } - return info - - -class LeRobotPolicy(Agent): - - def __init__( - self, - policy_name: str = "pi05", - default_checkpoint_path: str = "lerobot/pi05_base", - device: str = "cuda:0", - n_action_steps: int = 30, - temporal_ensemble_coeff: float | None = None, - rename_map: dict[str, str] | None = None, - **kwargs, - ) -> None: - super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) - - self.policy_name = policy_name - self.device = device - self.n_action_steps = n_action_steps - self.temporal_ensemble_coeff = temporal_ensemble_coeff - checkpoint_path = self.checkpoint_path or self.default_checkpoint_path - if self.checkpoint_step is not None: - checkpoint_path = checkpoint_path.format(checkpoint_step=self.checkpoint_step) - self.path = checkpoint_path - - if rename_map is not None: - self.rename_map = rename_map - else: - self.rename_map = {} - - # self.rename_map = { - # "head": "image", - # "left_wrist": "image2", - # "right_wrist": "image3", - # } - - def initialize(self): - from collections import deque - - import torch - from lerobot.policies.factory import get_policy_class, make_pre_post_processors - from torchvision.transforms import v2 - - # from vlagents import train_xvla - - self.policy = get_policy_class(self.policy_name).from_pretrained(self.path) - self.policy.config.n_action_steps = self.n_action_steps - - if self.policy_name == "act": - from lerobot.policies.act.modeling_act import ACTTemporalEnsembler - - if self.temporal_ensemble_coeff is not None: - self.policy.config.temporal_ensemble_coeff = self.temporal_ensemble_coeff - self.policy.temporal_ensembler = ACTTemporalEnsembler( - self.temporal_ensemble_coeff, - self.policy.config.chunk_size, - ) - elif hasattr(self.policy, "temporal_ensembler"): - delattr(self.policy, "temporal_ensembler") - - if self.policy.config.temporal_ensemble_coeff is None: - self.policy._action_queue = deque([], maxlen=self.policy.config.n_action_steps) - - self._expected_image_shapes = { - key.removeprefix("observation.images."): tuple(feature.shape) - for key, feature in self.policy.config.input_features.items() - if key.startswith("observation.images.") - } - self._camera_transforms = { - key: v2.Compose( - [ - v2.ToImage(), - v2.Resize((height, width)), - v2.ToDtype(torch.float32, scale=True), - v2.ToPureTensor(), - ] - ) - for key, (_, height, width) in self._expected_image_shapes.items() - } - # self.policy.config.device = self.device - self.policy.to(self.device) - self.policy.eval() - - preprocessor_overrides = { - "device_processor": {"device": self.device}, - # "rename_observations_processor": {"rename_map": self.rename_map}, - } - - self.preprocessor, self.postprocessor = make_pre_post_processors( - policy_cfg=self.policy.config, - pretrained_path=self.path, - preprocessor_overrides=preprocessor_overrides, - ) - - def act(self, obs: Obs) -> Act: - import torch - - super().act(obs) - - observation = { - "observation.state": torch.as_tensor(np.array(obs.state, copy=True)).to(torch.float32), - "task": self.instruction, - } - - for key, img_data in obs.cameras.items(): - expected_shape = self._expected_image_shapes.get(self.rename_map.get(key, key)) - assert expected_shape is not None - observation[f"observation.images.{self.rename_map.get(key, key)}"] = self._camera_transforms[ - self.rename_map.get(key, key) - ](np.array(img_data, copy=True)) - - observation = self.preprocessor(observation) - - with torch.inference_mode(): - action = self.policy.select_action(observation) - # action = self.policy.predict_action_chunk(observation) - action = self.postprocessor(action) - - if isinstance(action, torch.Tensor): - action = action.detach().float().cpu().numpy() - - action = np.squeeze(action, axis=0) - return Act(action=np.asarray(action, dtype=np.float32)) - - def reset(self, obs: Obs, instruction: Any, **kwargs) -> dict[str, Any]: - info = super().reset(obs, instruction, **kwargs) - self.policy.reset() - return info - - -class VjepaAC(Agent): - - def __init__( - self, - cfg_path: str, - model_name: str = "vjepa2_ac_vit_giant", - default_checkpoint_path: str = "", - **kwargs, - ) -> None: - super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) - import yaml - - self.cfg_path = cfg_path - with open(self.cfg_path, "r") as f: - self.cfg = yaml.safe_load(f) - - self.model_name = model_name - - def initialize(self): - # torch import - import torch - - # VJEPA imports - from app.vjepa_droid.transforms import make_transforms - from notebooks.utils.world_model_wrapper import WorldModel - - self.device = self.cfg.get("device", "cuda") - self.goal_img = self.cfg.get("goal_img", "exp_1.png") - - # data config - cfgs_data = self.cfg.get("data") - crop_size = cfgs_data.get("crop_size", 256) - - # data augs - cfgs_data_aug = self.cfg.get("data_aug") - use_aa = cfgs_data_aug.get("auto_augment", False) - horizontal_flip = cfgs_data_aug.get("horizontal_flip", False) - motion_shift = cfgs_data_aug.get("motion_shift", False) - ar_range = cfgs_data_aug.get("random_resize_aspect_ratio", [3 / 4, 4 / 3]) - rr_scale = cfgs_data_aug.get("random_resize_scale", [0.3, 1.0]) - reprob = cfgs_data_aug.get("reprob", 0.0) - - # cfgs_mpc_args config - cfgs_mpc_args = self.cfg.get("mpc_args") - self.rollout_horizon = cfgs_mpc_args.get("rollout_horizon", 2) - samples = cfgs_mpc_args.get("samples", 25) - topk = cfgs_mpc_args.get("topk", 10) - cem_steps = cfgs_mpc_args.get("cem_steps", 1) - momentum_mean = cfgs_mpc_args.get("momentum_mean", 0.15) - momentum_mean_gripper = cfgs_mpc_args.get("momentum_mean_gripper", 0.15) - momentum_std = cfgs_mpc_args.get("momentum_std", 0.75) - momentum_std_gripper = cfgs_mpc_args.get("momentum_std_gripper", 0.15) - maxnorm = cfgs_mpc_args.get("maxnorm", 0.075) - verbose = cfgs_mpc_args.get("verbose", True) - - # Initialize transform (random-resize-crop augmentations) - self.transform = make_transforms( - random_horizontal_flip=horizontal_flip, - random_resize_aspect_ratio=ar_range, - random_resize_scale=rr_scale, - reprob=reprob, - auto_augment=use_aa, - motion_shift=motion_shift, - crop_size=crop_size, - ) - - # load model - encoder, predictor = torch.hub.load( - "./", self.model_name, source="local", pretrained=True # root of the vjepa source code # model type - ) - - # load model to cuda - encoder.to(self.device) - predictor.to(self.device) - - # World model wrapper initialization - tokens_per_frame = int((crop_size // encoder.patch_size) ** 2) - self.world_model = WorldModel( - encoder=encoder, - predictor=predictor, - tokens_per_frame=tokens_per_frame, - mpc_args={ - "rollout": self.rollout_horizon, - "samples": samples, - "topk": topk, - "cem_steps": cem_steps, - "momentum_mean": momentum_mean, - "momentum_mean_gripper": momentum_mean_gripper, - "momentum_std": momentum_std, - "momentum_std_gripper": momentum_std_gripper, - "maxnorm": maxnorm, - "verbose": verbose, - }, - normalize_reps=True, - device=self.device, - ) - - def act(self, obs: Obs) -> Act: - # torch imports - import torch - from torchvision.io import decode_jpeg - - super().act(obs) - - with torch.no_grad(): - - # read from camera-stream - side = obs.cameras["rgb_side"] - - # [3, 720, 1280] -> [1, 720, 1280, 3] i.e, [T, C, Patches, dim] - side = torch.permute(side, (1, 2, 0)).unsqueeze(0) - - # [1, 720, 1280, 3] -> [1, 3, 1, 256, 1408] i.e, [B, C, T, Patches, dim] - input_image_tensor = (self.transform(side)[None, :]).to( - device=self.device, dtype=torch.float, non_blocking=True - ) - # Pre-trained VJEPA 2 ENCODER: [1, 3, 1, 256, 1408] -> [1, 256, 1408] - z_n = self.world_model.encode(input_image_tensor) - - # [1, 7] -> [B, state_dim] - # TODO: check gripper state convention - # in DROID: 0: is close to 0.86: is open? - # In rcs 0: is close and 1: is open - s_n = ( - torch.tensor((np.concatenate(([obs.info["xyzrpy"], [1 - obs.gripper]]), axis=0))) # [1-obs.gripper] - .unsqueeze(0) - .to(self.device, dtype=torch.float, non_blocking=True) - ) - - # Action conditioned predictor and zero-shot action inference with CEM - actions = self.world_model.infer_next_action(z_n, s_n, self.goal_rep) # [rollout_horizon, 7] - - first_action = actions[0].cpu() - first_action[-1] = 1 - first_action[-1] - - return Act(action=np.array(first_action)) - - def reset(self, obs: Obs, instruction: Any, **kwargs) -> dict[str, Any]: - super().reset(obs, instruction, **kwargs) - # imports - import torch - - img = Image.open(self.goal_img) - - # time dim exp - goal_image = np.expand_dims(np.array(img), axis=0) - # batch dim exp - goal_image_tensor = torch.tensor(self.transform(goal_image)[None, :]).to( - device=self.device, dtype=torch.float, non_blocking=True - ) - - with torch.no_grad(): - self.goal_rep = self.world_model.encode(goal_image_tensor) - - return {} - - -class OpenPiModel(Agent): - - def __init__( - self, - train_config_name: str = "pi0_droid", - default_checkpoint_path: str = "gs://openpi-assets/checkpoints/pi0_droid", - execution_horizon=20, - **kwargs, - ) -> None: - super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) - from openpi.training import config - - logging.info(f"checkpoint_path: {self.checkpoint_path}, checkpoint_step: {self.checkpoint_step}") - self.openpi_path = self.checkpoint_path.format(checkpoint_step=self.checkpoint_step) - - self.cfg = config.get_config(train_config_name) - self.execution_horizon = execution_horizon - - self.chunk_counter = self.execution_horizon - self._cached_action_chunk = None - - def initialize(self): - from openpi.policies import policy_config - from openpi.shared import download - - checkpoint_dir = download.maybe_download(self.openpi_path) - - # Create a trained policy. - self.policy = policy_config.create_trained_policy(self.cfg, checkpoint_dir) - - def act(self, obs: Obs) -> Act: - if self.chunk_counter < self.execution_horizon: - self.chunk_counter += 1 - return Act(action=self._cached_action_chunk[self.chunk_counter]) - - else: - self.chunk_counter = 0 - observation = {f"observation/{k}": np.copy(v).transpose(2, 0, 1) for k, v in obs.cameras.items()} - observation.update( - { - # openpi expects 0 as gripper open and 1 as closed - "observation/state": np.concatenate([obs.info["joints"], [1 - obs.gripper]]), - "prompt": self.instruction, - } - ) - action_chunk = self.policy.infer(observation)["actions"] - - # convert gripper action into vlagents format - action_chunk[:, -1] = 1 - action_chunk[:, -1] - self._cached_action_chunk = action_chunk - - return Act(action=action_chunk[0]) - - def reset(self, obs: Obs, instruction: Any): - super().reset(obs, instruction) - self.chunk_counter = self.execution_horizon - self._cached_action_chunk = None - return {} - - -class OpenVLAModel(Agent): - # === Utilities === - SYSTEM_PROMPT = ( - "A chat between a curious user and an artificial intelligence assistant. " - "The assistant gives helpful, detailed, and polite answers to the user's questions." - ) - - def __init__( - self, - attn_implementation: str, - device: str, - unnorm_key: str, - default_checkpoint_path: str = "openvla/openvla-7b", - **kwargs, - ) -> None: - super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) - self.unnorm_key = unnorm_key - logging.info(f"Using unnorm_key: {self.unnorm_key}") - self.attn_implementation = attn_implementation - if self.checkpoint_step is None or self.checkpoint_path is None: - self.openvla_path = self.default_checkpoint_path - logging.info(f"Using default checkpoint path: {self.openvla_path}") - if self.unnorm_key != "viola": - logging.warning( - "unnorm_key should be 'viola' when using default path, ignoring unnorm_key and setting it to 'viola'" - ) - self.unnorm_key = "viola" - else: - self.openvla_path = self.checkpoint_path.format(checkpoint_step=self.checkpoint_step) - logging.info( - f"Using custom checkpoint path: {self.openvla_path} with checkpoint step: {self.checkpoint_step}" - ) - self.device = device - self.attn_implementation = attn_implementation - - def initialize(self): - import torch - from transformers import AutoModelForVision2Seq, AutoProcessor - - self.device = torch.device(self.device) if torch.cuda.is_available() else torch.device("cpu") - - # Load VLA Model using HF AutoClasses - self.processor = AutoProcessor.from_pretrained(self.openvla_path, trust_remote_code=True) - self.vla = AutoModelForVision2Seq.from_pretrained( - self.openvla_path, - attn_implementation=self.attn_implementation, - torch_dtype=torch.bfloat16, - low_cpu_mem_usage=True, - trust_remote_code=True, - ).to(self.device) - print("==========================") - print(self.vla.norm_stats.keys()) - - # [Hacky] Load Dataset Statistics from Disk (if passing a path to a fine-tuned model) - if os.path.isdir(self.openvla_path): - with open(Path(self.openvla_path) / "dataset_statistics.json", "r") as f: - self.vla.norm_stats = json.load(f) - - def get_openvla_prompt(self, instruction: str, openvla_path: Union[str, Path]) -> str: - if "v01" in openvla_path: - return f"{self.SYSTEM_PROMPT} USER: What action should the robot take to {instruction.lower()}? ASSISTANT:" - else: - return f"In: What action should the robot take to {instruction.lower()}?\nOut:" - - def act(self, obs: Obs) -> Act: - # no batch dimension here - import torch - - super().act(obs) - # Parse payload components - assert obs.cameras["rgb_side"].shape == (256, 256, 3), "wrong shape, use lanczos" - image = obs.cameras["rgb_side"] - unnorm_key = self.unnorm_key - - # Run VLA Inference - prompt = self.get_openvla_prompt(self.instruction, self.openvla_path) - inputs = self.processor(prompt, Image.fromarray(image).convert("RGB")).to(self.device, dtype=torch.bfloat16) - # to use temperature use: do_sample=True, temperature=50.0 - action = self.vla.predict_action(**inputs, unnorm_key=unnorm_key, do_sample=False) - # unsqueeze to add horizon dimension - return Act(action=action[None]) - - -class OctoModel(Agent): - """ - This model is trained with a window size of 2, predicting 7 dimensional actions 4 steps into the future. - Observations and tasks conform to the following spec: - - Observations: { - image_primary: ('batch', 'history_window', 256, 256, 3), - image_wrist: ('batch', 'history_window', 128, 128, 3), - } - Tasks: { - image_primary: ('batch', 256, 256, 3), - image_wrist: ('batch', 128, 128, 3), - language_instruction: { - attention_mask: ('batch', 16), - input_ids: ('batch', 16), - }, - } - - At inference, you may pass in any subset of these observation and task keys, with a history window up to 2 timesteps. - """ - - def __init__( - self, - horizon: int = 2, - unnorm_key: list[str] | None = None, - default_checkpoint_path: str = "hf://rail-berkeley/octo-base-1.5", - **kwargs, - ) -> None: - # default window size is 2 in octo - # default unnorm is viola as it used the fr3 - super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) - self.horizon = horizon - if unnorm_key is None: - self.unnorm_key = [] - else: - self.unnorm_key = unnorm_key - - # log checkpoint path and step and kwargs - logging.info(f"checkpoint_path: {self.checkpoint_path}, checkpoint_step: {self.checkpoint_step}") - logging.info(f"horizon: {self.horizon}") - logging.info(f"unnorm_key: {self.unnorm_key}") - logging.info(f"kwargs: {kwargs}") - if self.checkpoint_path is None: - self.octo_path = self.default_checkpoint_path - logging.info(f"Using default checkpoint path: {self.octo_path}") - if self.checkpoint_step is not None: - logging.warning( - "checkpoint_step should be None when using default path, ignoring checkpoint_step and setting it to None" - ) - self.checkpoint_step = None - if self.unnorm_key != ["viola"]: - logging.warning( - "unnorm_key should be ['viola'] when using default path, ignoring unnorm_key and setting it to ['viola']" - ) - self.unnorm_key = ["viola"] - else: - self.octo_path = self.checkpoint_path - logging.info(f"Using custom checkpoint path: {self.octo_path} with checkpoint step: {self.checkpoint_step}") - logging.info(f"Using unnorm_key: {self.unnorm_key}") - - def initialize(self): - os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = ( - "false" # disable preallocation of memory in jax (might make it less efficient) - ) - from octo.model.octo_model import OctoModel as _OctoModel - from octo.utils.train_callbacks import supply_rng - - self.model = _OctoModel.load_pretrained(self.octo_path, self.checkpoint_step) - - window_size = self.model.example_batch["observation"]["timestep_pad_mask"].shape[1] - if window_size < self.horizon: - logging.warning( - f"Horizon {self.horizon} is greater than the model's window size {window_size} with which the model has been trained with. " - ) - - self.trained_obs = self.model.example_batch["observation"].keys() - - self.horizon = self.horizon - self.history = deque(maxlen=self.horizon) - self.num_obs = 0 - - logging.info("==========================") - logging.info(self.model.dataset_statistics.keys()) - self.policy_fn = supply_rng( - partial( - self.model.sample_actions, - unnormalization_statistics=reduce(getitem, self.unnorm_key, self.model.dataset_statistics)["action"], - ), - ) - self.task = None - - def act(self, obs: Obs) -> Act: - # from octo.model.octo_model import _verify_shapes - import jax - from octo.utils.gym_wrappers import stack_and_pad - - super().act(obs) - assert self.task is not None, "forgot reset?" - # _verify_shapes(obs, , self.model.example_batch["observation"]) - - self.num_obs += 1 - - # single image - assert obs.cameras["rgb_side"].shape == (256, 256, 3), "wrong shape, use lanczos" - obs = {"image_primary": obs.cameras["rgb_side"]} - - self.history.append(obs) - assert len(self.history) == self.horizon, "forgot reset?" - full_obs = stack_and_pad(self.history, self.num_obs) - - actions = self.policy_fn( - jax.tree_map( - lambda x: x[None], - full_obs, - ), - self.task, - ) - # remove the batch dimension (batch, horizon, action) - return Act(action=np.array(actions[0, :, :])) - - def reset(self, obs: Obs, instruction: Any): - super().reset(obs, instruction) - assert obs.cameras["rgb_side"].shape == (256, 256, 3), "wrong shape" - obs = {"image_primary": obs.cameras["rgb_side"]} - self.task = self.model.create_tasks(texts=[instruction]) - self.num_obs = 1 - self.history.extend([obs] * self.horizon) - return {} - - -class OctoActionDistribution(OctoModel): - """ - this model does not support history window - dont use self.step and self.episode as this model is not used sequentially - """ - - def __init__(self, **kwargs) -> None: - assert kwargs["horizon"] == 1, "horizon must be 1 for OctoActionDistribution" - super().__init__(**kwargs) - - def act(self, obs: Obs) -> Act: - """ - Args: - Obs: - cameras: - rgb_side: np.ndarray[tuple[BATCH, H, W, Literal[3]], np.dtype[np.int8]] - info: - num_samples: int - Return: - Act: - action: None - info: - means: np.ndarray[tuple[BATCH, 7], np.dtype[np.float32]] - stds: np.ndarray[tuple[BATCH, 7], np.dtype[np.float32]] - """ - import jax - import jax.numpy as jnp - - self._from_shared_memory(obs) - - batch_size = obs.cameras["rgb_side"].shape[0] - assert obs.cameras["rgb_side"].shape == (batch_size, 256, 256, 3), "wrong shape" - assert self.instruction is not None, "forgot reset?" - num_samples = obs.info.get("num_samples", 1) - - x = jnp.array(obs.cameras["rgb_side"]) # BATCH, H, W, 3 - # x_expanded = x[:, None, :, :, :] - x_expanded = jnp.expand_dims(x, 1) - x_tiled = jnp.tile(x_expanded, (1, num_samples, 1, 1, 1)) # Shape: [BATCH, N, H, W, 3] - x_duplicated = x_tiled.reshape(-1, x.shape[1], x.shape[2], x.shape[3]) # Shape: [BATCH*N, H, W, 3] - full_obs = { - "image_primary": jnp.expand_dims(x_duplicated, 1), - "timestep_pad_mask": np.ones((batch_size * num_samples, 1)), - } - # full_obs = stack_and_pad(x_duplicated, 1) - tasks = self.model.create_tasks(texts=[self.instruction] * batch_size * num_samples) - actions = self.policy_fn( - full_obs, - tasks, - ) - # actions: [num_samples x BATCH, 4, 7] - # remove the horizon dimension and reshape to [BATCH, num_samples, 7] - actions = actions[:, 0, :].reshape(batch_size, num_samples, 7) - stds = jnp.std(actions, axis=1) - means = jnp.mean(actions, axis=1) - - stds = np.asarray(stds) - means = np.asarray(means) - - return Act(action=None, info={"means": means, "stds": stds, "actions": np.asarray(actions)}) - - def reset(self, obs, instruction): - self.instruction = instruction - return {} - - -class OpenVLADistribution(OpenVLAModel): - - def act(self, obs: Obs) -> Act: - # no batch dimension here - import torch - - self._from_shared_memory(obs) - - assert self.instruction is not None, "forgot reset?" - self.step += 1 - batch_size = obs.cameras["rgb_side"].shape[0] - - # Parse payload components - images = obs.cameras["rgb_side"] - actions = [] - unnorm_key = self.unnorm_key - num_samples = obs.info.get("num_samples", 1) - - # time it - import time - - t1 = time.time() - # Run VLA Inference - prompt = self.get_openvla_prompt(self.instruction, self.openvla_path) - - x_expanded = np.expand_dims(images, 1) - x_tiled = np.tile(x_expanded, (1, num_samples, 1, 1, 1)) # Shape: [BATCH, N, H, W, 3] - x_duplicated = x_tiled.reshape( - -1, images.shape[1], images.shape[2], images.shape[3] - ) # Shape: [BATCH*N, H, W, 3] - - for image in x_duplicated: - inputs = self.processor(prompt, Image.fromarray(image).convert("RGB")).to(self.device, dtype=torch.bfloat16) - # to use temperature use: do_sample=True, temperature=50.0 - action = self.vla.predict_action(**inputs, unnorm_key=unnorm_key, do_sample=False) - actions.append(action) - - t2 = time.time() - logging.info(f"needed time for {len(actions)} was {t2-t1}s") - - # unsqueeze to add horizon dimension - actions = np.stack(actions).astype(np.float32) - actions = actions.reshape(batch_size, num_samples, 7) - means = np.mean(actions, axis=1).astype(np.float32) - stds = np.std(actions, axis=1).astype(np.float32) - return Act(action=None, info={"means": means, "stds": stds, "actions": actions}) - - -AGENTS = dict( - test=TestAgent, - octo=OctoModel, - lerobot=LeRobotPolicy, - openvla=OpenVLAModel, - octodist=OctoActionDistribution, - openvladist=OpenVLADistribution, - openpi=OpenPiModel, - vjepa=VjepaAC, -) diff --git a/src/vlagents/policies/__init__.py b/src/vlagents/policies/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/vlagents/policies/interface.py b/src/vlagents/policies/interface.py new file mode 100644 index 0000000..fc64add --- /dev/null +++ b/src/vlagents/policies/interface.py @@ -0,0 +1,235 @@ +import base64 +from dataclasses import dataclass, field +from multiprocessing import resource_tracker, shared_memory +from typing import Any + +import numpy as np +import simplejpeg + +from vlagents import register_agent + + +@dataclass(kw_only=True) +class SharedMemoryPayload: + shm_name: str + shape: tuple[int, ...] + dtype: str = "uint8" + + +class CameraDataType: + SHARED_MEMORY = "shared_memory" + JPEG_ENCODED = "jpeg_encoded" + RAW = "raw" + + +@dataclass(kw_only=True) +class SingleObs: + cameras: dict[str, np.ndarray | SharedMemoryPayload | str] = field(default_factory=dict) + camera_data_type: str = CameraDataType.RAW + gripper: float | None = None + joints: np.ndarray | None = None + # translation in m and rotation around x, y, z (roll, pitch, yaw) axes in radians + xyzrpy: np.ndarray | None = None + # translation in m and quaternion in (x, y, z, w) format + tquat: np.ndarray | None = None + info: dict[str, Any] = field(default_factory=dict) + + +@dataclass(kw_only=True) +class Obs: + # dictionary for multiple robot arms + obs: dict[str, SingleObs] = field(default_factory=dict) + language_instruction: str | None = None + goal_image: np.ndarray | SharedMemoryPayload | str | None = None + goal_image_data_type: str = CameraDataType.RAW + + +@dataclass(kw_only=True) +class SingleAct: + action: np.ndarray + gripper: float | None = None + done: bool = False + info: dict[str, Any] = field(default_factory=dict) + + +@dataclass(kw_only=True) +class Act: + # action chunk with dictionary for multiple robot arms + acts: list[dict[str, SingleAct]] = field(default_factory=list) + + +class Agent: + def __init__( + self, + default_checkpoint_path: str, + checkpoint_path: str | None = None, + checkpoint_step: int | None = None, + ) -> None: + self.checkpoint_step = checkpoint_step + self.default_checkpoint_path = default_checkpoint_path + self.checkpoint_path = checkpoint_path + self.instruction: str | None = None + self.step = -1 + self._shm: dict[str, shared_memory.SharedMemory] = {} + + def initialize(self): + # heavy initialization, e.g. loading models + pass + + def _decode_image_payload( + self, + payload: np.ndarray | SharedMemoryPayload | str, + data_type: str, + ) -> np.ndarray: + if data_type == CameraDataType.RAW: + assert isinstance(payload, np.ndarray) + return payload + if data_type == CameraDataType.SHARED_MEMORY: + assert isinstance(payload, SharedMemoryPayload) + if payload.shm_name not in self._shm: + self._shm[payload.shm_name] = shared_memory.SharedMemory(payload.shm_name) + shm = self._shm[payload.shm_name] + return np.ndarray(payload.shape, dtype=payload.dtype, buffer=shm.buf) + if data_type == CameraDataType.JPEG_ENCODED: + assert isinstance(payload, str) + return simplejpeg.decode_jpeg(base64.urlsafe_b64decode(payload)) + raise ValueError(f"Unsupported camera data type: {data_type}") + + def _to_numpy(self, obs: Obs) -> Obs: + """Decode camera payloads in-place for every robot and goal image.""" + for single_obs in obs.obs.values(): + single_obs.cameras = { + camera_name: self._decode_image_payload(camera_data, single_obs.camera_data_type) + for camera_name, camera_data in single_obs.cameras.items() + } + single_obs.camera_data_type = CameraDataType.RAW + + if obs.goal_image is not None: + obs.goal_image = self._decode_image_payload(obs.goal_image, obs.goal_image_data_type) + obs.goal_image_data_type = CameraDataType.RAW + return obs + + def _require_single_arm(self, obs: Obs) -> tuple[str, SingleObs]: + if len(obs.obs) != 1: + raise ValueError(f"{type(self).__name__} currently supports exactly one arm, got {list(obs.obs.keys())}") + robot_name, single_obs = next(iter(obs.obs.items())) + return robot_name, single_obs + + def _single_obs_state(self, single_obs: SingleObs, *, include_gripper: bool = True) -> np.ndarray: + state_parts: list[np.ndarray] = [] + if single_obs.joints is not None: + state_parts.append(np.asarray(single_obs.joints, dtype=np.float32)) + if include_gripper and single_obs.gripper is not None: + state_parts.append(np.asarray([single_obs.gripper], dtype=np.float32)) + if not state_parts: + raise ValueError(f"{type(self).__name__} requires joints and/or gripper in the observation") + return np.concatenate(state_parts) + + def _single_step_act( + self, + robot_name: str, + action: np.ndarray, + *, + gripper: float | None = None, + done: bool = False, + info: dict[str, Any] | None = None, + ) -> Act: + return Act( + acts=[ + { + robot_name: SingleAct( + action=np.asarray(action, dtype=np.float32), + gripper=None if gripper is None else float(gripper), + done=done, + info={} if info is None else info, + ) + } + ] + ) + + def _chunk_act( + self, + robot_name: str, + action_chunk: np.ndarray, + *, + grippers: np.ndarray | list[float] | None = None, + infos: list[dict[str, Any] | None] | None = None, + done: bool = False, + ) -> Act: + actions = np.asarray(action_chunk, dtype=np.float32) + if actions.ndim == 1: + actions = actions[None, :] + if grippers is None: + gripper_values = [None] * len(actions) + else: + gripper_array = np.asarray(grippers, dtype=np.float32).reshape(-1) + if len(gripper_array) != len(actions): + raise ValueError("grippers must have the same length as the action chunk") + gripper_values = [float(value) for value in gripper_array] + info_values = infos or [None] * len(actions) + if len(info_values) != len(actions): + raise ValueError("infos must have the same length as the action chunk") + return Act( + acts=[ + { + robot_name: SingleAct( + action=actions[idx], + gripper=gripper_values[idx], + done=done and idx == len(actions) - 1, + info={} if info_values[idx] is None else info_values[idx], + ) + } + for idx in range(len(actions)) + ] + ) + + def act(self, obs: Obs) -> Act: + self.instruction = obs.language_instruction + self.step += 1 + self._to_numpy(obs) + return Act(acts=[]) + + def __enter__(self): + pass + + def __exit__(self, *args, **kwargs): + self.close() + + def close(self, *args, **kwargs): + for shm in self._shm.values(): + shm.close() + resource_tracker.unregister(shm._name, "shared_memory") + self._shm = {} + + +class TestAgent(Agent): + def __init__(self, **kwargs) -> None: + super().__init__(default_checkpoint_path="", **kwargs) + self.i = 0 + + def act(self, obs: Obs) -> Act: + super().act(obs) + assert len(obs.obs) == 1, "TestAgent currently expects a single robot observation" + robot_name, robot_obs = next(iter(obs.obs.items())) + info = { + "shapes": {k: v.shape for k, v in robot_obs.cameras.items()}, + "dtype": {k: v.dtype.name for k, v in robot_obs.cameras.items()}, + "data": {k: v for k, v in robot_obs.cameras.items()}, + } + a = Act( + acts=[ + { + robot_name: SingleAct( + action=np.array([0, 0, 0, 0, 0, 0], dtype=np.float32), + gripper=float(self.i % 2), + done=False, + info=info, + ) + } + ] + ) + self.i += 1 + return a + + +register_agent("test", TestAgent) diff --git a/src/vlagents/policies/lerobot.py b/src/vlagents/policies/lerobot.py new file mode 100644 index 0000000..a92bd13 --- /dev/null +++ b/src/vlagents/policies/lerobot.py @@ -0,0 +1,129 @@ +import numpy as np + +from vlagents import register_agent +from vlagents.policies.interface import Act, Agent, Obs + + +class LeRobotPolicy(Agent): + def __init__( + self, + policy_name: str = "pi05", + default_checkpoint_path: str = "lerobot/pi05_base", + device: str = "cuda:0", + temporal_ensemble_coeff: float | None = None, + rename_map: dict[str, str] | None = None, + **kwargs, + ) -> None: + super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) + + self.policy_name = policy_name + self.device = device + self.temporal_ensemble_coeff = temporal_ensemble_coeff + checkpoint_path = self.checkpoint_path or self.default_checkpoint_path + if self.checkpoint_step is not None: + checkpoint_path = checkpoint_path.format(checkpoint_step=self.checkpoint_step) + self.path = checkpoint_path + + if rename_map is not None: + self.rename_map = rename_map + else: + self.rename_map = {} + + # self.rename_map = { + # "head": "image", + # "left_wrist": "image2", + # "right_wrist": "image3", + # } + + def initialize(self): + + import torch + from lerobot.policies.factory import get_policy_class, make_pre_post_processors + from torchvision.transforms import v2 + + # from vlagents import train_xvla + + self.policy = get_policy_class(self.policy_name).from_pretrained(self.path) + + if self.policy_name == "act": + from lerobot.policies.act.modeling_act import ACTTemporalEnsembler + + if self.temporal_ensemble_coeff is not None: + self.policy.config.temporal_ensemble_coeff = self.temporal_ensemble_coeff + self.policy.temporal_ensembler = ACTTemporalEnsembler( + self.temporal_ensemble_coeff, + self.policy.config.chunk_size, + ) + elif hasattr(self.policy, "temporal_ensembler"): + delattr(self.policy, "temporal_ensembler") + + self._expected_image_shapes = { + key.removeprefix("observation.images."): tuple(feature.shape) + for key, feature in self.policy.config.input_features.items() + if key.startswith("observation.images.") + } + self._camera_transforms = { + key: v2.Compose( + [ + v2.ToImage(), + v2.Resize((height, width)), + v2.ToDtype(torch.float32, scale=True), + v2.ToPureTensor(), + ] + ) + for key, (_, height, width) in self._expected_image_shapes.items() + } + # self.policy.config.device = self.device + self.policy.to(self.device) + self.policy.eval() + + preprocessor_overrides = { + "device_processor": {"device": self.device}, + # "rename_observations_processor": {"rename_map": self.rename_map}, + } + + self.preprocessor, self.postprocessor = make_pre_post_processors( + policy_cfg=self.policy.config, + pretrained_path=self.path, + preprocessor_overrides=preprocessor_overrides, + ) + + def act(self, obs: Obs) -> Act: + import torch + + super().act(obs) + robot_name, single_obs = self._require_single_arm(obs) + + observation = { + "observation.state": torch.as_tensor(np.array(self._single_obs_state(single_obs), copy=True)).to( + torch.float32 + ), + "task": obs.language_instruction, + } + + for key, img_data in single_obs.cameras.items(): + renamed_key = self.rename_map.get(key, key) + expected_shape = self._expected_image_shapes.get(renamed_key) + assert expected_shape is not None + observation[f"observation.images.{renamed_key}"] = self._camera_transforms[renamed_key]( + np.array(img_data, copy=True) + ) + + observation = self.preprocessor(observation) + + with torch.inference_mode(): + action = self.policy.predict_action_chunk(observation) + action = self.postprocessor(action) + + if isinstance(action, torch.Tensor): + action = action.detach().float().cpu().numpy() + + action_chunk = np.squeeze(action, axis=0) # remove batch dimension + if action_chunk.ndim == 1: + action_chunk = action_chunk[None, :] + if action_chunk.shape[-1] < 1: + raise ValueError("LeRobot action chunk must include a gripper dimension") + return self._chunk_act(robot_name, action_chunk[:, :-1], grippers=action_chunk[:, -1]) + + +register_agent("lerobot", LeRobotPolicy) diff --git a/src/vlagents/policies/octo.py b/src/vlagents/policies/octo.py new file mode 100644 index 0000000..ccca3bd --- /dev/null +++ b/src/vlagents/policies/octo.py @@ -0,0 +1,122 @@ +import logging +import os +from collections import deque +from functools import partial, reduce +from operator import getitem + +import numpy as np + +from vlagents import register_agent +from vlagents.policies.interface import Act, Agent, Obs + + +class OctoModel(Agent): + """ + This model is trained with a window size of 2, predicting 7 dimensional actions 4 steps into the future. + Observations and tasks conform to the following spec: + + Observations: { + image_primary: ('batch', 'history_window', 256, 256, 3), + image_wrist: ('batch', 'history_window', 128, 128, 3), + } + Tasks: { + image_primary: ('batch', 256, 256, 3), + image_wrist: ('batch', 128, 128, 3), + language_instruction: { + attention_mask: ('batch', 16), + input_ids: ('batch', 16), + }, + } + + At inference, you may pass in any subset of these observation and task keys, with a history window up to 2 timesteps. + """ + + def __init__( + self, + horizon: int = 2, + unnorm_key: list[str] | None = None, + default_checkpoint_path: str = "hf://rail-berkeley/octo-base-1.5", + **kwargs, + ) -> None: + # default window size is 2 in octo + # default unnorm is viola as it used the fr3 + super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) + self.horizon = horizon + if unnorm_key is None: + self.unnorm_key = [] + else: + self.unnorm_key = unnorm_key + + # log checkpoint path and step and kwargs + logging.info(f"checkpoint_path: {self.checkpoint_path}, checkpoint_step: {self.checkpoint_step}") + logging.info(f"horizon: {self.horizon}") + logging.info(f"unnorm_key: {self.unnorm_key}") + logging.info(f"kwargs: {kwargs}") + if self.checkpoint_path is None: + self.octo_path = self.default_checkpoint_path + logging.info(f"Using default checkpoint path: {self.octo_path}") + if self.checkpoint_step is not None: + logging.warning( + "checkpoint_step should be None when using default path, ignoring checkpoint_step and setting it to None" + ) + self.checkpoint_step = None + if self.unnorm_key != ["viola"]: + logging.warning( + "unnorm_key should be ['viola'] when using default path, ignoring unnorm_key and setting it to ['viola']" + ) + self.unnorm_key = ["viola"] + else: + self.octo_path = self.checkpoint_path + logging.info(f"Using custom checkpoint path: {self.octo_path} with checkpoint step: {self.checkpoint_step}") + logging.info(f"Using unnorm_key: {self.unnorm_key}") + + def initialize(self): + os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = ( + "false" # disable preallocation of memory in jax (might make it less efficient) + ) + from octo.model.octo_model import OctoModel as _OctoModel + from octo.utils.train_callbacks import supply_rng + + self.model = _OctoModel.load_pretrained(self.octo_path, self.checkpoint_step) + + window_size = self.model.example_batch["observation"]["timestep_pad_mask"].shape[1] + if window_size < self.horizon: + logging.warning( + f"Horizon {self.horizon} is greater than the model's window size {window_size} with which the model has been trained with. " + ) + + self.trained_obs = self.model.example_batch["observation"].keys() + + logging.info("==========================") + logging.info(self.model.dataset_statistics.keys()) + self.policy_fn = supply_rng( + partial( + self.model.sample_actions, + unnormalization_statistics=reduce(getitem, self.unnorm_key, self.model.dataset_statistics)["action"], + ), + ) + + def act(self, obs: Obs) -> Act: + import jax + from octo.utils.gym_wrappers import stack_and_pad + + super().act(obs) + robot_name, single_obs = self._require_single_arm(obs) + assert single_obs.cameras["rgb_side"].shape == (256, 256, 3), "wrong shape, use lanczos" + image_obs = {"image_primary": single_obs.cameras["rgb_side"]} + history = deque([image_obs] * self.horizon, maxlen=self.horizon) + full_obs = stack_and_pad(history, self.horizon) + task = self.model.create_tasks(texts=[obs.language_instruction or ""]) + + actions = self.policy_fn( + jax.tree_map( + lambda x: x[None], + full_obs, + ), + task, + ) + action_chunk = np.asarray(actions[0, :, :], dtype=np.float32) + return self._chunk_act(robot_name, action_chunk[:, :-1], grippers=action_chunk[:, -1]) + + +register_agent("octo", OctoModel) diff --git a/src/vlagents/policies/openpi.py b/src/vlagents/policies/openpi.py new file mode 100644 index 0000000..feb3eff --- /dev/null +++ b/src/vlagents/policies/openpi.py @@ -0,0 +1,55 @@ +import logging + +import numpy as np + +from vlagents import register_agent +from vlagents.policies.interface import Act, Agent, Obs + + +class OpenPiModel(Agent): + def __init__( + self, + train_config_name: str = "pi0_droid", + default_checkpoint_path: str = "gs://openpi-assets/checkpoints/pi0_droid", + **kwargs, + ) -> None: + super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) + from openpi.training import config + + logging.info(f"checkpoint_path: {self.checkpoint_path}, checkpoint_step: {self.checkpoint_step}") + self.openpi_path = self.checkpoint_path.format(checkpoint_step=self.checkpoint_step) + + self.cfg = config.get_config(train_config_name) + + def initialize(self): + from openpi.policies import policy_config + from openpi.shared import download + + checkpoint_dir = download.maybe_download(self.openpi_path) + + # Create a trained policy. + self.policy = policy_config.create_trained_policy(self.cfg, checkpoint_dir) + + def act(self, obs: Obs) -> Act: + super().act(obs) + robot_name, single_obs = self._require_single_arm(obs) + observation = { + # OpenPI expects channel-first images: [H, W, C] -> [C, H, W] + f"observation/{k}": np.copy(v).transpose(2, 0, 1) + for k, v in single_obs.cameras.items() + } + observation.update( + { + # openpi expects 0 as gripper open and 1 as closed + "observation/state": np.concatenate( + [np.asarray(single_obs.joints, dtype=np.float32), [1 - float(single_obs.gripper or 0.0)]] + ), + "prompt": obs.language_instruction, + } + ) + action_chunk = np.asarray(self.policy.infer(observation)["actions"], dtype=np.float32) + action_chunk[:, -1] = 1 - action_chunk[:, -1] + return self._chunk_act(robot_name, action_chunk[:, :-1], grippers=action_chunk[:, -1]) + + +register_agent("openpi", OpenPiModel) diff --git a/src/vlagents/policies/openvla.py b/src/vlagents/policies/openvla.py new file mode 100644 index 0000000..baab38a --- /dev/null +++ b/src/vlagents/policies/openvla.py @@ -0,0 +1,95 @@ +import json +import logging +import os +from pathlib import Path +from typing import Union + +import numpy as np +from PIL import Image + +from vlagents import register_agent +from vlagents.policies.interface import Act, Agent, Obs + + +class OpenVLAModel(Agent): + # === Utilities === + SYSTEM_PROMPT = ( + "A chat between a curious user and an artificial intelligence assistant. " + "The assistant gives helpful, detailed, and polite answers to the user's questions." + ) + + def __init__( + self, + attn_implementation: str, + device: str, + unnorm_key: str, + default_checkpoint_path: str = "openvla/openvla-7b", + **kwargs, + ) -> None: + super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) + self.unnorm_key = unnorm_key + logging.info(f"Using unnorm_key: {self.unnorm_key}") + self.attn_implementation = attn_implementation + if self.checkpoint_step is None or self.checkpoint_path is None: + self.openvla_path = self.default_checkpoint_path + logging.info(f"Using default checkpoint path: {self.openvla_path}") + if self.unnorm_key != "viola": + logging.warning( + "unnorm_key should be 'viola' when using default path, ignoring unnorm_key and setting it to 'viola'" + ) + self.unnorm_key = "viola" + else: + self.openvla_path = self.checkpoint_path.format(checkpoint_step=self.checkpoint_step) + logging.info( + f"Using custom checkpoint path: {self.openvla_path} with checkpoint step: {self.checkpoint_step}" + ) + self.device = device + self.attn_implementation = attn_implementation + + def initialize(self): + import torch + from transformers import AutoModelForVision2Seq, AutoProcessor + + self.device = torch.device(self.device) if torch.cuda.is_available() else torch.device("cpu") + + # Load VLA Model using HF AutoClasses + self.processor = AutoProcessor.from_pretrained(self.openvla_path, trust_remote_code=True) + self.vla = AutoModelForVision2Seq.from_pretrained( + self.openvla_path, + attn_implementation=self.attn_implementation, + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + trust_remote_code=True, + ).to(self.device) + print("==========================") + print(self.vla.norm_stats.keys()) + + # [Hacky] Load Dataset Statistics from Disk (if passing a path to a fine-tuned model) + if os.path.isdir(self.openvla_path): + with open(Path(self.openvla_path) / "dataset_statistics.json", "r") as f: + self.vla.norm_stats = json.load(f) + + def get_openvla_prompt(self, instruction: str, openvla_path: Union[str, Path]) -> str: + if "v01" in openvla_path: + return f"{self.SYSTEM_PROMPT} USER: What action should the robot take to {instruction.lower()}? ASSISTANT:" + else: + return f"In: What action should the robot take to {instruction.lower()}?\nOut:" + + def act(self, obs: Obs) -> Act: + # no batch dimension here + import torch + + super().act(obs) + robot_name, single_obs = self._require_single_arm(obs) + assert single_obs.cameras["rgb_side"].shape == (256, 256, 3), "wrong shape, use lanczos" + image = single_obs.cameras["rgb_side"] + unnorm_key = self.unnorm_key + + prompt = self.get_openvla_prompt(obs.language_instruction or "", self.openvla_path) + inputs = self.processor(prompt, Image.fromarray(image).convert("RGB")).to(self.device, dtype=torch.bfloat16) + action = np.asarray(self.vla.predict_action(**inputs, unnorm_key=unnorm_key, do_sample=False), dtype=np.float32) + # OpenVLA returns a single step with gripper in the last dimension. + return self._single_step_act(robot_name, action[:-1], gripper=float(action[-1])) + + +register_agent("openvla", OpenVLAModel) diff --git a/src/vlagents/policies/vjepa.py b/src/vlagents/policies/vjepa.py new file mode 100644 index 0000000..433e887 --- /dev/null +++ b/src/vlagents/policies/vjepa.py @@ -0,0 +1,154 @@ +import numpy as np +from PIL import Image + +from vlagents import register_agent +from vlagents.policies.interface import Act, Agent, Obs + + +class VjepaAC(Agent): + def __init__( + self, + cfg_path: str, + model_name: str = "vjepa2_ac_vit_giant", + default_checkpoint_path: str = "", + **kwargs, + ) -> None: + super().__init__(default_checkpoint_path=default_checkpoint_path, **kwargs) + import yaml + + self.cfg_path = cfg_path + with open(self.cfg_path, "r") as f: + self.cfg = yaml.safe_load(f) + + self.model_name = model_name + + def initialize(self): + # torch import + import torch + + # VJEPA imports + from app.vjepa_droid.transforms import make_transforms + from notebooks.utils.world_model_wrapper import WorldModel + + self.device = self.cfg.get("device", "cuda") + self.goal_img = self.cfg.get("goal_img", "exp_1.png") + + # data config + cfgs_data = self.cfg.get("data") + crop_size = cfgs_data.get("crop_size", 256) + + # data augs + cfgs_data_aug = self.cfg.get("data_aug") + use_aa = cfgs_data_aug.get("auto_augment", False) + horizontal_flip = cfgs_data_aug.get("horizontal_flip", False) + motion_shift = cfgs_data_aug.get("motion_shift", False) + ar_range = cfgs_data_aug.get("random_resize_aspect_ratio", [3 / 4, 4 / 3]) + rr_scale = cfgs_data_aug.get("random_resize_scale", [0.3, 1.0]) + reprob = cfgs_data_aug.get("reprob", 0.0) + + # cfgs_mpc_args config + cfgs_mpc_args = self.cfg.get("mpc_args") + self.rollout_horizon = cfgs_mpc_args.get("rollout_horizon", 2) + samples = cfgs_mpc_args.get("samples", 25) + topk = cfgs_mpc_args.get("topk", 10) + cem_steps = cfgs_mpc_args.get("cem_steps", 1) + momentum_mean = cfgs_mpc_args.get("momentum_mean", 0.15) + momentum_mean_gripper = cfgs_mpc_args.get("momentum_mean_gripper", 0.15) + momentum_std = cfgs_mpc_args.get("momentum_std", 0.75) + momentum_std_gripper = cfgs_mpc_args.get("momentum_std_gripper", 0.15) + maxnorm = cfgs_mpc_args.get("maxnorm", 0.075) + verbose = cfgs_mpc_args.get("verbose", True) + + # Initialize transform (random-resize-crop augmentations) + self.transform = make_transforms( + random_horizontal_flip=horizontal_flip, + random_resize_aspect_ratio=ar_range, + random_resize_scale=rr_scale, + reprob=reprob, + auto_augment=use_aa, + motion_shift=motion_shift, + crop_size=crop_size, + ) + + # load model + encoder, predictor = torch.hub.load( + "./", + self.model_name, + source="local", + pretrained=True, # root of the vjepa source code # model type + ) + + # load model to cuda + encoder.to(self.device) + predictor.to(self.device) + + # World model wrapper initialization + tokens_per_frame = int((crop_size // encoder.patch_size) ** 2) + self.world_model = WorldModel( + encoder=encoder, + predictor=predictor, + tokens_per_frame=tokens_per_frame, + mpc_args={ + "rollout": self.rollout_horizon, + "samples": samples, + "topk": topk, + "cem_steps": cem_steps, + "momentum_mean": momentum_mean, + "momentum_mean_gripper": momentum_mean_gripper, + "momentum_std": momentum_std, + "momentum_std_gripper": momentum_std_gripper, + "maxnorm": maxnorm, + "verbose": verbose, + }, + normalize_reps=True, + device=self.device, + ) + + img = Image.open(self.goal_img) + # [H, W, C] -> [T=1, H, W, C] + goal_image = np.expand_dims(np.array(img), axis=0) + # [T=1, H, W, C] -> [B=1, C, T, crop, patches] + goal_image_tensor = torch.tensor(self.transform(goal_image)[None, :]).to( + device=self.device, dtype=torch.float, non_blocking=True + ) + with torch.no_grad(): + self.goal_rep = self.world_model.encode(goal_image_tensor) + + def act(self, obs: Obs) -> Act: + # torch imports + import torch + + super().act(obs) + robot_name, single_obs = self._require_single_arm(obs) + + with torch.no_grad(): + side = single_obs.cameras["rgb_side"] + # [H, W, C] -> [T=1, H, W, C] + side = torch.permute(torch.as_tensor(side), (1, 2, 0)).unsqueeze(0) + # [T=1, H, W, C] -> [B=1, C, T, crop, patches] + input_image_tensor = (self.transform(side)[None, :]).to( + device=self.device, dtype=torch.float, non_blocking=True + ) + # encoder output shape: [B=1, num_tokens, dim] + z_n = self.world_model.encode(input_image_tensor) + + if single_obs.xyzrpy is None: + raise ValueError("VjepaAC requires xyzrpy in SingleObs.xyzrpy") + xyzrpy = np.asarray(single_obs.xyzrpy, dtype=np.float32) + gripper = float(single_obs.gripper if single_obs.gripper is not None else 0.0) + # [xyzrpy(6), gripper(1)] -> [B=1, state_dim] + s_n = ( + torch.tensor(np.concatenate((xyzrpy, [1 - gripper]), axis=0)) + .unsqueeze(0) + .to(self.device, dtype=torch.float, non_blocking=True) + ) + + # predicted action chunk: [rollout_horizon, action_dim] + actions = np.asarray(self.world_model.infer_next_action(z_n, s_n, self.goal_rep).cpu(), dtype=np.float32) + # VJEPA uses the opposite gripper convention from vlagents. + actions[:, -1] = 1 - actions[:, -1] + + return self._chunk_act(robot_name, actions[:, :-1], grippers=actions[:, -1]) + + +register_agent("vjepa_ac", VjepaAC) diff --git a/src/vlagents/server.py b/src/vlagents/server.py index 269ee64..1463018 100644 --- a/src/vlagents/server.py +++ b/src/vlagents/server.py @@ -10,7 +10,7 @@ import rpyc from vlagents.client import dataclass_from_dict -from vlagents.policies import Agent, CameraDataType, Obs, SharedMemoryPayload +from vlagents.policies.interface import Agent, CameraDataType, Obs, SharedMemoryPayload logging.basicConfig( format="%(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s", @@ -21,7 +21,6 @@ @rpyc.service class AgentService(rpyc.Service): - # TODO: think if we should identify the connection with the instance GIT_ID = "git_id_remote.txt" GIT_ID_SUBMODULES = "git_id_submodules_remote.txt" GIT_DIFF = "git_diff_remote.txt" @@ -51,26 +50,16 @@ def act(self, obs_bytes: bytes) -> str: assert self._is_initialized, "AgentService not initialized, wait until is_initialized is True" # action, done, info obs = typing.cast(Obs, dataclass_from_dict(Obs, json_numpy.loads(obs_bytes))) - if obs.camera_data_type == CameraDataType.SHARED_MEMORY: - obs.cameras = { - camera_name: dataclass_from_dict(SharedMemoryPayload, camera_data) - for camera_name, camera_data in obs.cameras.items() - } + for single_obs in obs.obs.values(): + if single_obs.camera_data_type == CameraDataType.SHARED_MEMORY: + single_obs.cameras = { + camera_name: dataclass_from_dict(SharedMemoryPayload, camera_data) + for camera_name, camera_data in single_obs.cameras.items() + } + if obs.goal_image_data_type == CameraDataType.SHARED_MEMORY and obs.goal_image is not None: + obs.goal_image = dataclass_from_dict(SharedMemoryPayload, obs.goal_image) return json_numpy.dumps(asdict(self.agent.act(obs))) - @rpyc.exposed - def reset(self, args: bytes) -> str: - assert self._is_initialized, "AgentService not initialized, wait until is_initialized is True" - # info - obs, instruction, kwargs = json_numpy.loads(args) - obs_dclass = typing.cast(Obs, dataclass_from_dict(Obs, obs)) - if obs_dclass.camera_data_type == CameraDataType.SHARED_MEMORY: - obs_dclass.cameras = { - camera_name: dataclass_from_dict(SharedMemoryPayload, camera_data) - for camera_name, camera_data in obs_dclass.cameras.items() - } - return json_numpy.dumps(self.agent.reset(obs_dclass, instruction, **kwargs)) - @rpyc.exposed def name(self) -> str: return self._name @@ -81,7 +70,6 @@ def is_initialized(self) -> bool: @rpyc.exposed def git_status(self) -> str: - # TODO: put git commit hash and git diff into temp file and read it into string and send it over with TemporaryDirectory() as tmp_dir: # git commit has id os.system(f'git log --format="%H" -n 1 > {os.path.join(tmp_dir, self.GIT_ID)}') diff --git a/src/vlagents/train_xvla.py b/src/vlagents/train_xvla.py deleted file mode 100644 index 53425bd..0000000 --- a/src/vlagents/train_xvla.py +++ /dev/null @@ -1,53 +0,0 @@ -import torch -import torch.nn as nn -from lerobot.policies.xvla.action_hub import BaseActionSpace, register_action - -XVLA_DOMAIN_ID = 20 - - -@register_action("frankaduo") -class FrankaDuoActionSpace(BaseActionSpace): - """Custom action space for dual Franka setup.""" - - dim_action = 20 - - # Use lists for safe PyTorch advanced indexing - gripper_idx = (7, 15) - # All indices EXCEPT 7 and 15 - joint_idx = [0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14] - - def __init__(self): - super().__init__() - self.mse = nn.MSELoss() - self.bce = nn.BCEWithLogitsLoss() - - def compute_loss(self, pred, target): - """Define your loss computation.""" - # Corrected: Now computes MSE for BOTH Robot 1 and Robot 2 joints - joints_loss = self.mse(pred[..., self.joint_idx], target[..., self.joint_idx]) - - # Computes BCE for both grippers - gripper_loss = self.bce(pred[..., self.gripper_idx], target[..., self.gripper_idx]) - - return { - "joints_loss": joints_loss, - "gripper_loss": gripper_loss, - } - - def preprocess(self, proprio, action, mode="train"): - """Preprocess actions before training.""" - proprio_m = proprio.clone() - action_m = action.clone() if action is not None else None - - # Zero out both grippers - proprio_m[..., self.gripper_idx] = 0.0 - if action_m is not None: - action_m[..., self.gripper_idx] = 0.0 - - return proprio_m, action_m - - def postprocess(self, action): - """Post-process predictions for deployment.""" - # Apply sigmoid to both gripper logits - action[..., self.gripper_idx] = torch.sigmoid(action[..., self.gripper_idx]) - return action[..., :16] diff --git a/src/vlagents/wrappers.py b/src/vlagents/wrappers.py deleted file mode 100644 index 1da31c7..0000000 --- a/src/vlagents/wrappers.py +++ /dev/null @@ -1,32 +0,0 @@ -import gymnasium as gym - - -class HumanCameraWrapper(gym.ObservationWrapper): - """ - Flattens the rgbd mode observations into a dictionary with two keys, "rgbd" and "state" - - Args: - rgb (bool): Whether to include rgb images in the observation - depth (bool): Whether to include depth images in the observation - state (bool): Whether to include state data in the observation - - Note that the returned observations will have a "rgbd" or "rgb" or "depth" key depending on the rgb/depth bool flags. - """ - - def __init__(self, env) -> None: - self.base_env = env.unwrapped - super().__init__(env) - new_obs = self.observation(self.base_env._init_raw_obs) - self.base_env.update_obs_space(new_obs) - - def observation(self, observation: dict): - # ret = dict() - if not hasattr(self.env, "_has_reset") or not self.env._has_reset: - self.env.reset() - self.env._has_reset = True - # observation["sensor_data"]["human_camera"] = dict(rgb=self.env.render()) - observation["sensor_data"]["base_camera"] = dict(rgb=self.env.render()) - return observation - - -WRAPPERS = {"HumanCameraWrapper": HumanCameraWrapper}