|
|
| """
|
| Test MrBalance directly from Hugging Face Hub.
|
|
|
| Requirements:
|
| pip install -U torch transformers safetensors mujoco
|
|
|
| Files:
|
| balance_plate_rl.py
|
| test_mrbalance_hf.py
|
|
|
| Usage:
|
| python test_mrbalance_hf.py
|
| python test_mrbalance_hf.py --object sphere
|
| python test_mrbalance_hf.py --object egg
|
| python test_mrbalance_hf.py --object heavy_ball
|
|
|
| The script:
|
| 1. Downloads MrBalance from Hugging Face.
|
| 2. Loads it through AutoModel with trust_remote_code=True.
|
| 3. Creates the original MuJoCo environment.
|
| 4. Uses the Hugging Face policy to control the plate.
|
| 5. Prints episode statistics.
|
| 6. Optionally launches the MuJoCo viewer.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import argparse
|
| import time
|
|
|
| import numpy as np
|
| import torch
|
| from transformers import AutoModel
|
|
|
| from balance_plate_rl import (
|
| DEVICE,
|
| MAX_EPISODE_STEPS,
|
| OBJECT_TYPES,
|
| PlateBalanceEnv,
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
| MODEL_ID = "fromziro/MrBalance"
|
|
|
| DEFAULT_EPISODES = 5
|
| DEFAULT_OBJECT = "sphere"
|
|
|
| DEVICE_TO_USE = torch.device(
|
| DEVICE
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
| def load_model():
|
| print("=" * 72)
|
| print("Loading MrBalance from Hugging Face")
|
| print("=" * 72)
|
|
|
| print(f"Model: {MODEL_ID}")
|
| print(f"Device: {DEVICE_TO_USE}")
|
| print()
|
|
|
| model = AutoModel.from_pretrained(
|
| MODEL_ID,
|
| trust_remote_code=True,
|
| )
|
|
|
| model = model.to(DEVICE_TO_USE)
|
| model.eval()
|
|
|
| print("[load] Model loaded successfully.")
|
| print()
|
|
|
|
|
| print(
|
| f"[model] Observation size: "
|
| f"{model.config.observation_size}"
|
| )
|
|
|
| print(
|
| f"[model] Hidden size: "
|
| f"{model.config.hidden_size}"
|
| )
|
|
|
| print(
|
| f"[model] Intermediate size: "
|
| f"{model.config.intermediate_size}"
|
| )
|
|
|
| print(
|
| f"[model] Bottleneck size: "
|
| f"{model.config.bottleneck_size}"
|
| )
|
|
|
| print(
|
| f"[model] Action size: "
|
| f"{model.config.action_size}"
|
| )
|
|
|
| print()
|
|
|
| return model
|
|
|
|
|
|
|
|
|
|
|
|
|
| def run_episode(
|
| model,
|
| object_name: str,
|
| seed: int,
|
| render: bool = False,
|
| ):
|
| env = PlateBalanceEnv(
|
| seed=seed,
|
| render=render,
|
| active_objects=[object_name],
|
| )
|
|
|
| obs = env.reset(
|
| specific_object=object_name
|
| )
|
|
|
| total_reward = 0.0
|
| episode_length = 0
|
|
|
| if render:
|
| try:
|
| import mujoco.viewer
|
|
|
| viewer_context = (
|
| mujoco.viewer.launch_passive(
|
| env.model,
|
| env.data,
|
| )
|
| )
|
| except Exception as exc:
|
| env.close()
|
| raise RuntimeError(
|
| f"Could not launch MuJoCo viewer: {exc}"
|
| ) from exc
|
| else:
|
| viewer_context = None
|
|
|
| try:
|
| if viewer_context is not None:
|
| with viewer_context as viewer:
|
|
|
| while viewer.is_running():
|
|
|
| step_start = time.time()
|
|
|
| obs_tensor = torch.as_tensor(
|
| obs,
|
| dtype=torch.float32,
|
| device=DEVICE_TO_USE,
|
| ).unsqueeze(0)
|
|
|
| with torch.no_grad():
|
| output = model(
|
| obs_tensor,
|
| deterministic=True,
|
| )
|
|
|
| action = (
|
| output.action[0]
|
| .detach()
|
| .cpu()
|
| .numpy()
|
| .astype(np.float64)
|
| )
|
|
|
| obs, reward, terminated, truncated, info = (
|
| env.step(action)
|
| )
|
|
|
| total_reward += reward
|
| episode_length += 1
|
|
|
| viewer.sync()
|
|
|
| if terminated or truncated:
|
| break
|
|
|
|
|
| target_dt = 0.01
|
| elapsed = time.time() - step_start
|
|
|
| if elapsed < target_dt:
|
| time.sleep(
|
| target_dt - elapsed
|
| )
|
|
|
| else:
|
| while True:
|
|
|
| obs_tensor = torch.as_tensor(
|
| obs,
|
| dtype=torch.float32,
|
| device=DEVICE_TO_USE,
|
| ).unsqueeze(0)
|
|
|
| with torch.no_grad():
|
| output = model(
|
| obs_tensor,
|
| deterministic=True,
|
| )
|
|
|
| action = (
|
| output.action[0]
|
| .detach()
|
| .cpu()
|
| .numpy()
|
| .astype(np.float64)
|
| )
|
|
|
| obs, reward, terminated, truncated, info = (
|
| env.step(action)
|
| )
|
|
|
| total_reward += reward
|
| episode_length += 1
|
|
|
| if terminated or truncated:
|
| break
|
|
|
| finally:
|
| env.close()
|
|
|
| success = (
|
| not info["fallen"]
|
| and episode_length >= MAX_EPISODE_STEPS
|
| )
|
|
|
| return {
|
| "object": object_name,
|
| "reward": float(total_reward),
|
| "length": int(episode_length),
|
| "success": bool(success),
|
| "fallen": bool(info["fallen"]),
|
| "distance": float(info["distance"]),
|
| "velocity": float(info["object_velocity"]),
|
| }
|
|
|
|
|
|
|
|
|
|
|
|
|
| def sanity_check(model):
|
| """
|
| Verify that the Hugging Face model can actually execute inference
|
| independently of MuJoCo.
|
| """
|
|
|
| print("=" * 72)
|
| print("Running model sanity check")
|
| print("=" * 72)
|
|
|
| x = torch.randn(
|
| 4,
|
| 64,
|
| dtype=torch.float32,
|
| device=DEVICE_TO_USE,
|
| )
|
|
|
| with torch.no_grad():
|
| output = model(
|
| x,
|
| deterministic=True,
|
| )
|
|
|
| print(
|
| "[sanity] action shape:",
|
| tuple(output.action.shape),
|
| )
|
|
|
| print(
|
| "[sanity] value shape:",
|
| tuple(output.value.shape),
|
| )
|
|
|
| print(
|
| "[sanity] action range:",
|
| float(output.action.min()),
|
| "to",
|
| float(output.action.max()),
|
| )
|
|
|
| print(
|
| "[sanity] example action:",
|
| output.action[0].detach().cpu().numpy(),
|
| )
|
|
|
| print("[sanity] PASS")
|
| print()
|
|
|
|
|
|
|
|
|
|
|
|
|
| def main():
|
| parser = argparse.ArgumentParser(
|
| description="Test MrBalance from Hugging Face."
|
| )
|
|
|
| parser.add_argument(
|
| "--object",
|
| type=str,
|
| default=DEFAULT_OBJECT,
|
| choices=OBJECT_TYPES,
|
| help="Object to balance.",
|
| )
|
|
|
| parser.add_argument(
|
| "--episodes",
|
| type=int,
|
| default=DEFAULT_EPISODES,
|
| help="Number of evaluation episodes.",
|
| )
|
|
|
| parser.add_argument(
|
| "--render",
|
| action="store_true",
|
| help="Launch interactive MuJoCo viewer.",
|
| )
|
|
|
| parser.add_argument(
|
| "--seed",
|
| type=int,
|
| default=12345,
|
| help="Evaluation seed.",
|
| )
|
|
|
| args = parser.parse_args()
|
|
|
| model = load_model()
|
|
|
| sanity_check(model)
|
|
|
| print("=" * 72)
|
| print(
|
| f"Testing object: {args.object}"
|
| )
|
| print(
|
| f"Episodes: {args.episodes}"
|
| )
|
| print("=" * 72)
|
| print()
|
|
|
| results = []
|
|
|
| for episode in range(args.episodes):
|
|
|
| print(
|
| f"[episode {episode + 1}/{args.episodes}] "
|
| f"Running..."
|
| )
|
|
|
| result = run_episode(
|
| model=model,
|
| object_name=args.object,
|
| seed=args.seed + episode,
|
| render=args.render,
|
| )
|
|
|
| results.append(result)
|
|
|
| status = (
|
| "SUCCESS"
|
| if result["success"]
|
| else (
|
| "FALL"
|
| if result["fallen"]
|
| else "TIMEOUT"
|
| )
|
| )
|
|
|
| print(
|
| f" status: {status}"
|
| )
|
|
|
| print(
|
| f" reward: {result['reward']:.2f}"
|
| )
|
|
|
| print(
|
| f" length: {result['length']}"
|
| )
|
|
|
| print(
|
| f" distance: {result['distance']:.4f} m"
|
| )
|
|
|
| print(
|
| f" velocity: {result['velocity']:.4f} m/s"
|
| )
|
|
|
| print()
|
|
|
|
|
|
|
|
|
|
|
| rewards = [
|
| r["reward"]
|
| for r in results
|
| ]
|
|
|
| lengths = [
|
| r["length"]
|
| for r in results
|
| ]
|
|
|
| distances = [
|
| r["distance"]
|
| for r in results
|
| ]
|
|
|
| successes = sum(
|
| r["success"]
|
| for r in results
|
| )
|
|
|
| print("=" * 72)
|
| print("RESULTS")
|
| print("=" * 72)
|
|
|
| print(
|
| f"Object: {args.object}"
|
| )
|
|
|
| print(
|
| f"Reward mean: {np.mean(rewards):.2f}"
|
| )
|
|
|
| print(
|
| f"Reward std: {np.std(rewards):.2f}"
|
| )
|
|
|
| print(
|
| f"Length mean: {np.mean(lengths):.1f}"
|
| )
|
|
|
| print(
|
| f"Survival rate: "
|
| f"{100.0 * sum(not r['fallen'] for r in results) / len(results):.1f}%"
|
| )
|
|
|
| print(
|
| f"Success rate: "
|
| f"{100.0 * successes / len(results):.1f}%"
|
| )
|
|
|
| print(
|
| f"Tracking error: "
|
| f"{np.mean(distances):.4f} m"
|
| )
|
|
|
| print("=" * 72)
|
|
|
|
|
| if __name__ == "__main__":
|
| main() |
