Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,197 @@
|
||||
"""In-process driver: DirectModelHandle + N x SubprocEnvHandle.
|
||||
|
||||
The model lives in the driver process while simulator envs live in subprocesses.
|
||||
Evaluation proceeds in chunk-level lockstep:
|
||||
1. Seed the driver process and construct the model handle.
|
||||
2. Start one subprocess env handle per worker.
|
||||
3. Claim a task-local frame from JobState.
|
||||
4. Reset envs, batch active observations, run model.predict_batch, then
|
||||
fan out action chunks to env subprocesses.
|
||||
5. Complete all episodes in the frame and move to the next frame.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
from wall_x._vendor.harrix.eval_config import EvalConfig
|
||||
from wall_x._vendor.harrix.drivers.inproc.env_handle import SubprocEnvHandle
|
||||
from wall_x._vendor.harrix.drivers.inproc.model_handle import DirectModelHandle
|
||||
from wall_x._vendor.harrix.drivers.job_state import JobState
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def run(cfg: EvalConfig) -> None:
|
||||
import wall_x._vendor.harrix.envs # noqa: F401 trigger env register
|
||||
import wall_x._vendor.harrix.adapters # noqa: F401 trigger adapter register
|
||||
from wall_x._vendor.harrix.envs.registry import enumerate_episodes_for
|
||||
from wall_x._vendor.harrix.utils.seed import set_seed_everywhere
|
||||
|
||||
# 1) Seed the driver process.
|
||||
set_seed_everywhere(cfg.env.seed)
|
||||
|
||||
# 2) Build DirectModelHandle; the model is loaded in the driver process.
|
||||
logger.info("Constructing DirectModelHandle in the driver process")
|
||||
t_model = time.time()
|
||||
model_handle = DirectModelHandle(cfg)
|
||||
logger.info("Model loaded in %.1fs", time.time() - t_model)
|
||||
|
||||
# 3) Start env subprocesses.
|
||||
logger.info("Starting %s SubprocEnvHandle(s)", cfg.runtime.num_workers)
|
||||
env_handles = [
|
||||
SubprocEnvHandle(cfg, worker_id=i) for i in range(cfg.runtime.num_workers)
|
||||
]
|
||||
|
||||
# 4) JobState runs in frame-sync mode for lockstep evaluation.
|
||||
os.makedirs(cfg.runtime.log_dir, exist_ok=True)
|
||||
log_path = os.path.join(cfg.runtime.log_dir, "state.jsonl")
|
||||
report_path = os.path.join(cfg.runtime.log_dir, "report.json")
|
||||
episodes = enumerate_episodes_for(cfg)
|
||||
logger.info("env.type=%r; scheduled %s episodes", cfg.env.type, len(episodes))
|
||||
state = JobState(
|
||||
episodes,
|
||||
log_path,
|
||||
batch_sync_mode=True,
|
||||
batch_size=cfg.runtime.num_workers,
|
||||
)
|
||||
|
||||
# 5) Main loop, one frame at a time.
|
||||
t0 = time.time()
|
||||
frame_idx = 0
|
||||
while not state.is_drained():
|
||||
frame_eps = state.claim_frame()
|
||||
if not frame_eps:
|
||||
# Previous frame still has in-flight episodes.
|
||||
time.sleep(0.05)
|
||||
continue
|
||||
|
||||
results = _run_frame(frame_eps, model_handle, env_handles, cfg, frame_idx)
|
||||
for ep, res in zip(frame_eps, results):
|
||||
if "_error" in res:
|
||||
state.fail(ep, str(res["_error"]))
|
||||
else:
|
||||
state.complete(ep, res)
|
||||
frame_idx += 1
|
||||
|
||||
elapsed = time.time() - t0
|
||||
logger.info("All episodes finished in %.1fs (%.1f min)", elapsed, elapsed / 60)
|
||||
|
||||
state.dump_final(report_path)
|
||||
with open(report_path) as f:
|
||||
report = json.load(f)
|
||||
overall = report["overall"]
|
||||
logger.info(
|
||||
"attempted=%s, successes=%s, success_rate=%.2f%%, failed=%s",
|
||||
overall["attempted"],
|
||||
overall["successes"],
|
||||
overall["success_rate"] * 100,
|
||||
overall["failed"],
|
||||
)
|
||||
logger.info("Report: %s", report_path)
|
||||
logger.info("State log: %s", log_path)
|
||||
|
||||
for h in env_handles:
|
||||
h.shutdown()
|
||||
model_handle.shutdown()
|
||||
|
||||
|
||||
def _run_frame(
|
||||
frame_eps: list,
|
||||
model_handle: DirectModelHandle,
|
||||
env_handles: list,
|
||||
cfg: EvalConfig,
|
||||
frame_idx: int,
|
||||
) -> list[dict]:
|
||||
"""Run one task-local frame with chunk-level lockstep.
|
||||
|
||||
Active workers are batched together at each chunk boundary. Workers that
|
||||
already finished no longer participate in later forwards.
|
||||
"""
|
||||
from wall_x._vendor.harrix.envs.libero_common import encode_raw_obs
|
||||
|
||||
n = len(frame_eps)
|
||||
t_frame_start = time.time()
|
||||
|
||||
# ---- a) reset: fan out, then gather ----
|
||||
for i in range(n):
|
||||
env_handles[i].submit_reset(tuple(frame_eps[i]))
|
||||
initials = [env_handles[i].wait_reset() for i in range(n)]
|
||||
|
||||
obs_list = [r["obs"] for r in initials]
|
||||
instr_list = [r["instruction"] for r in initials]
|
||||
status = [
|
||||
{
|
||||
"done": False,
|
||||
"success": False,
|
||||
"steps": 0,
|
||||
"task_desc": initials[i].get("task_desc", ""),
|
||||
}
|
||||
for i in range(n)
|
||||
]
|
||||
|
||||
max_rounds = cfg.env.libero.max_infer_times
|
||||
|
||||
# ---- b) chunk lockstep ----
|
||||
for round_idx in range(max_rounds):
|
||||
active = [i for i in range(n) if not status[i]["done"]]
|
||||
if not active:
|
||||
break
|
||||
|
||||
payloads = [
|
||||
{
|
||||
"observation": encode_raw_obs(obs_list[i]),
|
||||
"instruction": instr_list[i],
|
||||
"noise": None,
|
||||
}
|
||||
for i in active
|
||||
]
|
||||
chunks = model_handle.predict_batch(payloads)
|
||||
|
||||
# fan-out submit
|
||||
for k, i in enumerate(active):
|
||||
env_handles[i].submit_execute_chunk(chunks[k])
|
||||
# gather
|
||||
for k, i in enumerate(active):
|
||||
try:
|
||||
r = env_handles[i].wait_execute_chunk()
|
||||
except Exception as e:
|
||||
status[i]["_error"] = str(e)
|
||||
status[i]["done"] = True
|
||||
continue
|
||||
obs_list[i] = r["obs"]
|
||||
status[i]["steps"] += r["steps"]
|
||||
if r["done"]:
|
||||
status[i]["done"] = True
|
||||
status[i]["success"] = True
|
||||
|
||||
for i in range(n):
|
||||
env_handles[i].submit_finalize_episode(status[i]["success"])
|
||||
for i in range(n):
|
||||
env_handles[i].wait_finalize_episode()
|
||||
|
||||
elapsed_frame = time.time() - t_frame_start
|
||||
logger.info(
|
||||
"frame=%s n=%s succ=%s/%s elapsed=%.1fs",
|
||||
frame_idx,
|
||||
n,
|
||||
sum(s["success"] for s in status),
|
||||
n,
|
||||
elapsed_frame,
|
||||
)
|
||||
|
||||
return [
|
||||
{
|
||||
"success": s["success"],
|
||||
"steps": s["steps"],
|
||||
"elapsed_sec": round(elapsed_frame / max(1, n), 3),
|
||||
"task_desc": s["task_desc"],
|
||||
**({"_error": s["_error"]} if "_error" in s else {}),
|
||||
}
|
||||
for s in status
|
||||
]
|
||||
@@ -0,0 +1,152 @@
|
||||
"""Subprocess env handle used by the in-process driver.
|
||||
|
||||
The model remains in the driver process. Each env subprocess receives reset and
|
||||
execute-chunk commands through a pipe.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _subproc_main(cfg, worker_id, child_conn):
|
||||
"""Env subprocess entry point."""
|
||||
# Spawned children inherit env vars, but set these explicitly for launchers
|
||||
# that did not configure them. Robosuite validates the EGL id against the
|
||||
# CUDA_VISIBLE_DEVICES environment string, so keep the same visible id here.
|
||||
cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") or "0"
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
|
||||
# EGL device index is always 0 after CUDA_VISIBLE_DEVICES remapping.
|
||||
os.environ["MUJOCO_EGL_DEVICE_ID"] = os.environ.get("MUJOCO_EGL_DEVICE_ID") or "0"
|
||||
|
||||
# Env workers run robosuite only. Do not import torch here; otherwise many
|
||||
# subprocesses may initialize CUDA contexts and compete with the driver model.
|
||||
seed = cfg.env.seed + worker_id
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
os.environ["PYTHONHASHSEED"] = str(seed)
|
||||
|
||||
import wall_x._vendor.harrix.envs # noqa: F401 trigger register
|
||||
from wall_x._vendor.harrix.envs.registry import build_env
|
||||
|
||||
env = build_env(cfg, worker_id)
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
cmd, args = child_conn.recv()
|
||||
except EOFError:
|
||||
break
|
||||
try:
|
||||
if cmd == "reset_episode":
|
||||
result = env.reset_episode(tuple(args))
|
||||
elif cmd == "execute_chunk":
|
||||
result = env.execute_chunk(args)
|
||||
elif cmd == "finalize_episode":
|
||||
env.finalize_episode(bool(args))
|
||||
result = None
|
||||
elif cmd == "shutdown":
|
||||
env.shutdown()
|
||||
child_conn.send(("ok", None))
|
||||
break
|
||||
else:
|
||||
raise ValueError(f"unknown cmd {cmd!r}")
|
||||
child_conn.send(("ok", result))
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
child_conn.send(("err", f"{e}\n{traceback.format_exc()}"))
|
||||
finally:
|
||||
try:
|
||||
child_conn.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
class SubprocEnvHandle:
|
||||
"""Synchronous pipe wrapper around one env subprocess.
|
||||
|
||||
Usage:
|
||||
h.submit_reset(ep_id)
|
||||
...
|
||||
h.wait_reset() # -> {"obs", "instruction", "task_desc"}
|
||||
h.submit_execute_chunk(actions)
|
||||
...
|
||||
h.wait_execute_chunk() # -> {"obs", "done", "steps"}
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, worker_id: int):
|
||||
# Use spawn instead of fork because the driver may already hold a CUDA
|
||||
# context for the model.
|
||||
ctx = mp.get_context("spawn")
|
||||
self._parent_conn, child_conn = ctx.Pipe()
|
||||
self._proc = ctx.Process(
|
||||
target=_subproc_main,
|
||||
args=(cfg, worker_id, child_conn),
|
||||
daemon=False,
|
||||
name=f"infer-subproc-w{worker_id}",
|
||||
)
|
||||
self._proc.start()
|
||||
child_conn.close()
|
||||
self._has_pending = False
|
||||
|
||||
def _send(self, cmd: str, args):
|
||||
self._parent_conn.send((cmd, args))
|
||||
self._has_pending = True
|
||||
|
||||
def _recv(self):
|
||||
if not self._has_pending:
|
||||
raise RuntimeError("no pending request to wait for")
|
||||
status, payload = self._parent_conn.recv()
|
||||
self._has_pending = False
|
||||
if status == "err":
|
||||
raise RuntimeError(f"subproc env error (w={self._proc.name}): {payload}")
|
||||
return payload
|
||||
|
||||
def submit_reset(self, ep_id):
|
||||
self._send("reset_episode", ep_id)
|
||||
|
||||
def wait_reset(self):
|
||||
return self._recv()
|
||||
|
||||
def submit_execute_chunk(self, actions):
|
||||
self._send("execute_chunk", np.asarray(actions, dtype=np.float32))
|
||||
|
||||
def wait_execute_chunk(self):
|
||||
return self._recv()
|
||||
|
||||
def submit_finalize_episode(self, success: bool):
|
||||
self._send("finalize_episode", success)
|
||||
|
||||
def wait_finalize_episode(self):
|
||||
return self._recv()
|
||||
|
||||
def finalize_episode(self, success: bool):
|
||||
"""Blocking convenience wrapper."""
|
||||
self.submit_finalize_episode(success)
|
||||
return self.wait_finalize_episode()
|
||||
|
||||
def reset_episode(self, ep_id):
|
||||
"""Blocking convenience wrapper."""
|
||||
self.submit_reset(ep_id)
|
||||
return self.wait_reset()
|
||||
|
||||
def execute_chunk(self, actions):
|
||||
"""Blocking convenience wrapper."""
|
||||
self.submit_execute_chunk(actions)
|
||||
return self.wait_execute_chunk()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
try:
|
||||
self._send("shutdown", None)
|
||||
self._recv()
|
||||
except Exception:
|
||||
pass
|
||||
if self._proc.is_alive():
|
||||
self._proc.join(timeout=5)
|
||||
if self._proc.is_alive():
|
||||
self._proc.terminate()
|
||||
self._proc.join(timeout=2)
|
||||
@@ -0,0 +1,25 @@
|
||||
"""In-process model handle.
|
||||
|
||||
The model is constructed in the driver process and calls the adapter directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class DirectModelHandle:
|
||||
def __init__(self, cfg):
|
||||
# Trigger adapter registration.
|
||||
import wall_x._vendor.harrix.adapters # noqa: F401
|
||||
from wall_x._vendor.harrix.adapters.registry import build_adapter
|
||||
|
||||
self._adapter = build_adapter(cfg)
|
||||
|
||||
@property
|
||||
def chunk_horizon(self) -> int:
|
||||
return self._adapter.chunk_horizon
|
||||
|
||||
def predict_batch(self, payloads):
|
||||
return self._adapter.predict_batch(payloads)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
pass
|
||||
Reference in New Issue
Block a user