File size: 2,322 Bytes
977b06e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | import json
import os
from pathlib import Path
import numpy as np
from simpler_env.evaluation.maniskill2_evaluator import maniskill2_evaluator
# from IPython import embed; embed()
from examples.SimplerEnv.eval_files.bridgert_ik_trace import install_ik_trace
from examples.SimplerEnv.eval_files.custom_argparse import get_args
from examples.SimplerEnv.eval_files.model2simpler_interface import ModelClient
def start_debugpy_once():
import debugpy
if getattr(start_debugpy_once, "_started", False):
return
debugpy.listen(("0.0.0.0", 10092))
print("🔍 Waiting for VSCode attach on 0.0.0.0:10092 ...")
debugpy.wait_for_client()
start_debugpy_once._started = True
if __name__ == "__main__":
args = get_args()
os.environ["DISPLAY"] = ""
# prevent a single jax process from taking up all the GPU memory
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
install_ik_trace()
if os.getenv("DEBUG", False):
start_debugpy_once()
model = ModelClient(
policy_ckpt_path=args.ckpt_path, # to get unnormalization stats
policy_setup=args.policy_setup,
port=args.port,
action_scale=args.action_scale,
cfg_scale=1.5, # cfg from 1.5 to 7 also performs well
)
# policy model creation; update this if you are using a new policy model
# run real-to-sim evaluation
success_arr = maniskill2_evaluator(model, args)
print(args)
print(" " * 10, "Average success", np.mean(success_arr))
if args.result_json:
result = {
"schema_version": 2,
"env_name": args.env_name,
"episode_start": args.obj_episode_range[0],
"episode_end": args.obj_episode_range[1],
"episode_indices": list(range(*args.obj_episode_range)),
"checkpoint": str(Path(args.ckpt_path).resolve()),
"server_metadata": getattr(model, "server_metadata", {}),
"successes": np.asarray(success_arr, dtype=bool).tolist(),
}
destination = Path(args.result_json)
destination.parent.mkdir(parents=True, exist_ok=True)
temporary = destination.with_suffix(destination.suffix + ".tmp")
temporary.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
os.replace(temporary, destination)
|