56 lines
1.3 KiB
Python
56 lines
1.3 KiB
Python
import dataclasses
|
|
import logging
|
|
import pathlib
|
|
|
|
import env as _env
|
|
from openpi_client import action_chunk_broker
|
|
from openpi_client import websocket_client_policy as _websocket_client_policy
|
|
from openpi_client.runtime import runtime as _runtime
|
|
from openpi_client.runtime.agents import policy_agent as _policy_agent
|
|
import saver as _saver
|
|
import tyro
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Args:
|
|
out_dir: pathlib.Path = pathlib.Path("data/aloha_sim/videos")
|
|
|
|
task: str = "gym_aloha/AlohaTransferCube-v0"
|
|
seed: int = 0
|
|
|
|
action_horizon: int = 10
|
|
|
|
host: str = "0.0.0.0"
|
|
port: int = 8000
|
|
|
|
display: bool = False
|
|
|
|
|
|
def main(args: Args) -> None:
|
|
runtime = _runtime.Runtime(
|
|
environment=_env.AlohaSimEnvironment(
|
|
task=args.task,
|
|
seed=args.seed,
|
|
),
|
|
agent=_policy_agent.PolicyAgent(
|
|
policy=action_chunk_broker.ActionChunkBroker(
|
|
policy=_websocket_client_policy.WebsocketClientPolicy(
|
|
host=args.host,
|
|
port=args.port,
|
|
),
|
|
action_horizon=args.action_horizon,
|
|
)
|
|
),
|
|
subscribers=[
|
|
_saver.VideoSaver(args.out_dir),
|
|
],
|
|
max_hz=50,
|
|
)
|
|
|
|
runtime.run()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO, force=True)
|
|
tyro.cli(main)
|