vlagents 0.0.1__py3-none-any.whl

tests/test_connection.py

@@ -0,0 +1,80 @@
+import subprocess
+from time import sleep
+
+import numpy as np
+
+from vlagents.client import RemoteAgent
+from vlagents.evaluator_envs import start_server
+from vlagents.policies import Act, Obs
+
+
+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()
+
+    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
+
+
+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()
+
+
+def test_connection_numpy_serialization():
+    with start_server("test", {}, 8080, "localhost") as p:
+        sleep(2)
+        agent = RemoteAgent("localhost", 8080, "test")
+        with agent:
+            while not agent.is_initialized():
+                sleep(0.1)
+            _test_connection(agent)
+        p.send_signal(subprocess.signal.SIGINT)
+
+
+def test_connection_numpy_shm():
+    with start_server("test", {}, 8080, "localhost") as p:
+        sleep(2)
+        agent = RemoteAgent("localhost", 8080, "test", on_same_machine=True)
+        with agent:
+            while not agent.is_initialized():
+                sleep(0.1)
+            _test_connection(agent)
+        p.send_signal(subprocess.signal.SIGINT)
+
+
+def test_connection_numpy_jpeg():
+    with start_server("test", {}, 8080, "localhost") as p:
+        sleep(2)
+        agent = RemoteAgent("localhost", 8080, "test", jpeg_encoding=True)
+        with agent:
+            while not agent.is_initialized():
+                sleep(0.1)
+            _test_connection_jpeg(agent)
+        p.send_signal(subprocess.signal.SIGINT)

vlagents/__init__.py

@@ -0,0 +1 @@
+__all__ = ["__doc__", "__version__", "policies", "client", "server"]

vlagents/__main__.py

@@ -0,0 +1,400 @@
+import copy
+import json
+import logging
+import os
+from multiprocessing import Pool
+from pathlib import Path
+from typing import Annotated
+
+import numpy as np
+import rpyc
+import typer
+import wandb
+
+# to use non-inline backend, necessary for the case
+# 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.server import AgentService
+
+main_app = typer.Typer(help="CLI tool for the vlagents library.")
+
+
+def wandb_log_git_diff(path: str):
+    path = Path(path)
+    git_path = path / "git"
+    git_path.mkdir(parents=True, exist_ok=True)
+    log_git_diff(git_path)
+    wandb.log_artifact(git_path, type="directory", name="git")
+
+
+def log_git_diff(path: str):
+    # git id
+    git_id = os.path.join(path, "git_id.txt")
+    os.system(f'git log --format="%H" -n 1 > {git_id}')
+
+    # submodule git ids
+    git_submodules = os.path.join(path, "git_submodules.txt")
+    os.system(f"git submodule status > {git_submodules}")
+
+    # get git diff
+    git_diff = os.path.join(path, "git_diff.txt")
+    os.system(f"git diff --submodule=diff > {git_diff}")
+
+
+@main_app.command()
+def start_server(
+    agent_name: Annotated[str, typer.Argument(help="Agent name to run.")],
+    kwargs: Annotated[str, typer.Option(help="args to start the agent.")] = "{}",
+    port: Annotated[int, typer.Option(help="Port to run the server on.")] = 8080,
+    host: Annotated[str, typer.Option(help="Host to run the server on.")] = "localhost",
+):
+    """Runs eval server."""
+    agent = AGENTS[agent_name](**json.loads(kwargs))
+    service = AgentService(agent, agent_name)
+    with service:
+        t = rpyc.ThreadedServer(
+            service, port=port, hostname=host, protocol_config={"allow_pickle": True, "allow_public_attrs": True}
+        )
+        t.start()
+
+
+def _per_process(
+    args: tuple[int, AgentConfig, list[EvalConfig], int, int | None, int],
+) -> tuple[np.ndarray, list[list[list[float]]], list[float], int]:
+    step, _agent_cfg, eval_cfgs, episodes, n_processes, nth_gpu = args
+    logging.info(f"Starting evaluation for step {step}")
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(nth_gpu)
+    os.environ["CAM_PATH"] = f"{os.environ['RUN_PATH']}/videos/{step}"
+    agent_cfg = copy.deepcopy(_agent_cfg)
+    agent_cfg.agent_kwargs["checkpoint_step"] = step
+
+    per_env_results_last_reward, per_env_results_rewards = evaluation(
+        agent_cfg=agent_cfg, eval_cfgs=eval_cfgs, episodes=episodes, n_processes=n_processes
+    )
+    logging.info(f"Finished evaluation for step {step}")
+    flatten_rewards = [[item for sublist in env_rewards for item in sublist] for env_rewards in per_env_results_rewards]
+    mean_rewards = [np.mean(env_rewards) if env_rewards else 0.0 for env_rewards in flatten_rewards]
+    logging.info("Returning results for step %s", step)
+    return per_env_results_last_reward, per_env_results_rewards, mean_rewards, step
+
+
+@main_app.command()
+def run_eval_post_training(
+    wandb_project: Annotated[str, typer.Option(help="weights and biases logging project.")],
+    wandb_entity: Annotated[str, typer.Option(help="weights and biases logging entity.")],
+    wandb_note: Annotated[str, typer.Option(help="weights and biases logging note.")],
+    wandb_name: Annotated[str, typer.Option(help="weights and biases logging name.")],
+    output_path: Annotated[str, typer.Option(help="Path to store the run results.")],
+    wandb_group: Annotated[str | None, typer.Option(help="weights and biases logging name.")] = None,
+    steps: Annotated[str | None, typer.Option(help="steps to evaluate.")] = None,
+    episodes: Annotated[int, typer.Option(help="Number of episodes to run.")] = 100,
+    n_processes: Annotated[int | None, typer.Option(help="Number of processes to run.")] = None,
+    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": {}}]',
+    agent_cfg: Annotated[
+        str, typer.Option(help="Agent configuration.")
+    ] = '{"host": "localhost", "port": 8080, "agent_name": "Test", "agent_kwargs": {}, "python_path": "python"}',
+):
+    """
+    post training eval which goes over all checkpoints
+    - each checkpoint with many envs
+    """
+    if steps is None:
+        steps = [None]
+    else:
+        steps = json.loads(steps)
+
+    if wandb_group == "":
+        wandb_group = None
+
+    wandb.init(
+        entity=wandb_entity,
+        resume="allow",
+        project=wandb_project,
+        # config=dict(agent_name=agent_name, agent_kwargs=json.loads(kwargs), eval_cfgs=json.loads(eval_cfgs)),
+        notes=wandb_note,
+        job_type="eval",
+        name=wandb_name,
+        group=wandb_group,
+    )
+    wandb_log_git_diff(output_path)
+    wandb.run.log_code(".")
+
+    wandb.define_metric(
+        "total/success",
+        step_metric="train_step",
+        overwrite=False,
+        step_sync=False,
+        hidden=False,
+        summary="max",
+    )
+    wandb.define_metric(
+        "total/last_step_reward",
+        step_metric="train_step",
+        overwrite=False,
+        step_sync=False,
+        hidden=False,
+        summary="max",
+    )
+    wandb.define_metric(
+        "total/total_steps",
+        step_metric="train_step",
+        overwrite=False,
+        step_sync=False,
+        hidden=False,
+        summary="min",
+    )
+    wandb.define_metric(
+        "total/mean_reward",
+        step_metric="train_step",
+        overwrite=False,
+        step_sync=False,
+        hidden=False,
+        summary="max",
+    )
+    eval_cfgs = [EvalConfig(**cfg) for cfg in json.loads(eval_cfgs)]
+    for idx, env in enumerate(eval_cfgs):
+        wandb.define_metric(
+            f"{env.env_id}/success",
+            step_metric="train_step",
+            overwrite=False,
+            step_sync=False,
+            hidden=False,
+            summary="max",
+        )
+        wandb.define_metric(
+            f"{env.env_id}/last_step_reward",
+            step_metric="train_step",
+            overwrite=False,
+            step_sync=False,
+            hidden=False,
+            summary="max",
+        )
+        wandb.define_metric(
+            f"{env.env_id}/total_steps",
+            step_metric="train_step",
+            overwrite=False,
+            step_sync=False,
+            hidden=False,
+            summary="min",
+        )
+        wandb.define_metric(
+            f"{env.env_id}/mean_reward",
+            step_metric="train_step",
+            overwrite=False,
+            step_sync=False,
+            hidden=False,
+            summary="max",
+        )
+
+    # distribute gpus equally
+    gpus_ids = [i % n_gpus for i in range(len(steps))]
+
+    # spawn n processes and run in parallel
+
+    agent_cfgs = [AgentConfig(**json.loads(agent_cfg)) for _ in steps]
+    for idx in range(len(steps)):
+        agent_cfgs[idx].port += idx
+    with Pool(n_processes) as p:
+        args = [(step, agent_cfgs[idx], eval_cfgs, episodes, 1, gpus_ids[idx]) for idx, step in enumerate(steps)]
+        results = p.map(_per_process, args)
+    logging.info("Finished evaluation")
+
+    for result in results:
+        per_env_results_last_reward, per_env_results_rewards, mean_rewards, step = result
+        step = step if step is not None else 0
+        wandb_log_dict = {
+            "total/success": per_env_results_last_reward.mean(axis=(0, 1))[0],
+            "total/last_step_reward": per_env_results_last_reward.mean(axis=(0, 1))[1],
+            "total/total_steps": per_env_results_last_reward.mean(axis=(0, 1))[2],
+            "total/mean_reward": np.mean(mean_rewards),
+            "train_step": step,
+        }
+        # log for each env
+        for idx, env in enumerate(eval_cfgs):
+            wandb_log_dict.update(
+                {
+                    f"{env.env_id}/success": per_env_results_last_reward[idx].mean(axis=0)[0],
+                    f"{env.env_id}/last_step_reward": per_env_results_last_reward[idx].mean(axis=0)[1],
+                    f"{env.env_id}/total_steps": per_env_results_last_reward[idx].mean(axis=0)[2],
+                    f"{env.env_id}/mean_reward": mean_rewards[idx],
+                }
+            )
+        wandb.log(wandb_log_dict, step=step, commit=True)
+
+        path = write_results(
+            per_env_results_last_reward,
+            per_env_results_rewards,
+            eval_cfgs,
+            agent_cfg=agent_cfgs[0],
+            out=output_path,
+        )
+        wandb.log_artifact(path, type="file", name="results", aliases=[f"step_{step}"])
+
+
+@main_app.command()
+def run_eval_during_training(
+    wandb_id: Annotated[str, typer.Option(help="weights and biases logging id.")],
+    wandb_group: Annotated[str, typer.Option(help="weights and biases logging group.")],
+    wandb_project: Annotated[str, typer.Option(help="weights and biases logging project.")],
+    wandb_entity: Annotated[str, typer.Option(help="weights and biases logging entity.")],
+    wandb_note: Annotated[str, typer.Option(help="weights and biases logging note.")],
+    wandb_name: Annotated[str, typer.Option(help="weights and biases logging name.")],
+    output_path: Annotated[str, typer.Option(help="Path to store the run results.")],
+    wandb_first: Annotated[bool, typer.Option(help="whether its the first eval.")] = False,
+    episodes: Annotated[int, typer.Option(help="Number of episodes to run.")] = 100,
+    n_processes: Annotated[int | None, typer.Option(help="Number of processes to run.")] = None,
+    eval_cfgs: Annotated[
+        str, typer.Option(help="Evaluation configurations.")
+    ] = '[{"env": "rcs/SimplePickUpSim-v0", "kwargs": {}}]',
+    agent_cfg: Annotated[
+        str, typer.Option(help="Agent configuration.")
+    ] = '{"host": "localhost", "port": 8080, "agent_name": "Test", "agent_kwargs": {}, "python_path": "python"}',
+):
+    """
+    during training eval, all need to use the same id
+    - just for one model, but many envs
+    - can be new run but at least in the same project and same group as the training
+    """
+    assert (
+        agent_cfg["agent_name"] != "Test"
+    ), "agent_cfg needs to be passed as a json argument. See the default for an example."
+
+    if wandb_first:
+        wandb.init(
+            id=wandb_id,
+            entity=wandb_entity,
+            resume="allow",
+            group=wandb_group,
+            project=wandb_project,
+            # config=dict(agent_name=agent_name, agent_kwargs=json.loads(kwargs), eval_cfgs=json.loads(eval_cfgs)),
+            notes=wandb_note,
+            job_type="eval",
+            name=wandb_name,
+        )
+
+        wandb_log_git_diff(output_path)
+        wandb.run.log_code(".")
+    else:
+        wandb.init(id=wandb_id, entity=wandb_entity, resume="must", project=wandb_project)
+
+    eval_cfgs = [EvalConfig(**cfg) for cfg in json.loads(eval_cfgs)]
+
+    agent_cfg = AgentConfig(**json.loads(agent_cfg))
+    step = agent_cfg.agent_kwargs.get("checkpoint_step", 0)
+    step = step if step is not None else 0
+
+    per_env_results_last_reward, per_env_results_rewards = evaluation(
+        agent_cfg=agent_cfg, eval_cfgs=eval_cfgs, episodes=episodes, n_processes=n_processes
+    )
+
+    # return is [envs, episodes, 3(success, reward, steps)], [envs, episodes, rewards for all steps in the episode]
+
+    flatten_rewards = [[item for sublist in env_rewards for item in sublist] for env_rewards in per_env_results_rewards]
+    mean_rewards = [np.mean(env_rewards) if env_rewards else 0.0 for env_rewards in flatten_rewards]
+
+    # these new define metric to also not work with several jobs
+    # wandb says that logging can only be done in run groups
+    if wandb_first:
+        wandb.define_metric(
+            "total/success",
+            step_metric="train_step",
+            overwrite=False,
+            step_sync=False,
+            hidden=False,
+            summary="max",
+        )
+        wandb.define_metric(
+            "total/last_step_reward",
+            step_metric="train_step",
+            overwrite=False,
+            step_sync=False,
+            hidden=False,
+            summary="max",
+        )
+        wandb.define_metric(
+            "total/total_steps",
+            step_metric="train_step",
+            overwrite=False,
+            step_sync=False,
+            hidden=False,
+            summary="min",
+        )
+        wandb.define_metric(
+            "total/mean_reward",
+            step_metric="train_step",
+            overwrite=False,
+            step_sync=False,
+            hidden=False,
+            summary="max",
+        )
+        for idx, env in enumerate(eval_cfgs):
+            wandb.define_metric(
+                f"{env.env_id}/success",
+                step_metric="train_step",
+                overwrite=False,
+                step_sync=False,
+                hidden=False,
+                summary="max",
+            )
+            wandb.define_metric(
+                f"{env.env_id}/last_step_reward",
+                step_metric="train_step",
+                overwrite=False,
+                step_sync=False,
+                hidden=False,
+                summary="max",
+            )
+            wandb.define_metric(
+                f"{env.env_id}/total_steps",
+                step_metric="train_step",
+                overwrite=False,
+                step_sync=False,
+                hidden=False,
+                summary="min",
+            )
+            wandb.define_metric(
+                f"{env.env_id}/mean_reward",
+                step_metric="train_step",
+                overwrite=False,
+                step_sync=False,
+                hidden=False,
+                summary="max",
+            )
+
+    wandb_log_dict = {
+        "total/success": per_env_results_last_reward.mean(axis=(0, 1))[0],
+        "total/last_step_reward": per_env_results_last_reward.mean(axis=(0, 1))[1],
+        "total/total_steps": per_env_results_last_reward.mean(axis=(0, 1))[2],
+        "total/mean_reward": np.mean(mean_rewards),
+        "train_step": step,
+    }
+
+    # log for each env
+    for idx, env in enumerate(eval_cfgs):
+        wandb_log_dict.update(
+            {
+                f"{env.env_id}/success": per_env_results_last_reward[idx].mean(axis=0)[0],
+                f"{env.env_id}/last_step_reward": per_env_results_last_reward[idx].mean(axis=0)[1],
+                f"{env.env_id}/total_steps": per_env_results_last_reward[idx].mean(axis=0)[2],
+                f"{env.env_id}/mean_reward": mean_rewards[idx],
+            }
+        )
+    wandb.log(wandb_log_dict, step=step, commit=True)
+    path = write_results(
+        per_env_results_last_reward,
+        per_env_results_rewards,
+        eval_cfgs,
+        agent_cfg=agent_cfg,
+        out=output_path,
+    )
+    wandb.log_artifact(path, type="file", name="results", aliases=[f"step_{step}"])
+
+
+if __name__ == "__main__":
+    main_app()

vlagents/client.py

@@ -0,0 +1,107 @@
+import base64
+import dataclasses
+from dataclasses import asdict
+from multiprocessing import shared_memory
+from typing import Any
+
+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:
+        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
+
+
+class RemoteAgent(Agent):
+    def __init__(self, host: str, port: int, model: str, on_same_machine: bool = False, jpeg_encoding: bool = False):
+        """Connect to a remote agent service.
+
+        Args:
+            host (str): Hostname or IP address of the remote agent service.
+            port (int): Port number of the remote agent service.
+            model (str): Name of the model to connect to.
+            on_same_machine (bool, optional): If True, assumes the agent is running on the same machine and uses
+                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.
+        """
+        self.on_same_machine = on_same_machine
+        self.jpeg_encoding = jpeg_encoding
+        self._shm: dict[str, shared_memory.SharedMemory] = {}
+        self.c = rpyc.connect(
+            host, port, config={"allow_pickle": True, "allow_public_attrs": True, "sync_request_timeout": 300}
+        )
+        assert model == self.c.root.name()
+
+    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
+        return obs
+
+    def act(self, obs: Obs) -> Act:
+        obs = self._process(obs)
+        obs = json_numpy.dumps(asdict(obs))
+        # action, done, info
+        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
+        return json_numpy.loads(self.c.root.reset(json_numpy.dumps((obs_dict, instruction, kwargs))))
+
+    def git_status(self) -> str:
+        return json_numpy.loads(self.c.root.git_status())
+
+    def is_initialized(self) -> bool:
+        return self.c.root.is_initialized()
+
+    def close(self):
+        for shm in self._shm.values():
+            shm.close()
+            shm.unlink()
+        self._shm = {}
+        self.c.close()
+
+
+if __name__ == "__main__":
+    # to test the connection
+    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)
+    print(agent.act(obs))
+    print(agent.act(obs))