hanabhi/gridworld-env
0
1#!/usr/bin/env python32# Copyright (c) Meta Platforms, Inc. and affiliates.3# All rights reserved.4#5# This source code is licensed under the BSD-style license found in the6# LICENSE file in the root directory of this source tree.7 8"""9Unity ML-Agents Environment Example Usage10 11This script demonstrates how to use the Unity ML-Agents environment12through the OpenEnv interface, with support for direct mode, server mode,13and Docker-based deployment.14 15=============================================================================16USAGE EXAMPLES (run from the OpenEnv repository root)17=============================================================================18 191. DIRECT MODE (Recommended for quick testing - no server required)20 ----------------------------------------------------------------21 Runs the Unity environment directly in-process.22 This is the simplest way to get started.23 24 # Run with graphics (default: 1280x720 window)25 python examples/unity_simple.py --direct26 27 # Run with custom window size28 python examples/unity_simple.py --direct --width 1920 --height 108029 30 # Run headless (no graphics, faster for training)31 python examples/unity_simple.py --direct --no-graphics --time-scale 2032 33 # Run 3DBall environment for 5 episodes34 python examples/unity_simple.py --direct --env 3DBall --episodes 535 36 # Run alternating between PushBlock and 3DBall37 python examples/unity_simple.py --direct --env both --episodes 638 39 402. SERVER MODE (For client-server architecture)41 ---------------------------------------------42 First, start the server in one terminal, then connect with this script.43 44 Step 1: Start the server (in Terminal 1):45 cd envs/unity_env46 uvicorn server.app:app --host 0.0.0.0 --port 800047 48 Or with environment variables for custom settings:49 UNITY_WIDTH=1920 UNITY_HEIGHT=1080 uvicorn server.app:app --port 800050 UNITY_NO_GRAPHICS=1 UNITY_TIME_SCALE=20 uvicorn server.app:app --port 800051 52 Step 2: Run this script (in Terminal 2, from repo root):53 python examples/unity_simple.py --url http://localhost:800054 python examples/unity_simple.py --url http://localhost:8000 --env 3DBall --episodes 555 56 573. DOCKER MODE (For containerized deployment)58 -------------------------------------------59 Automatically starts a Docker container and connects to it.60 61 First, build the Docker image:62 cd envs/unity_env63 docker build -f server/Dockerfile -t unity-env:latest .64 65 Then run (from repo root):66 python examples/unity_simple.py --docker67 python examples/unity_simple.py --docker --width 1280 --height 72068 python examples/unity_simple.py --docker --no-graphics --time-scale 2069 python examples/unity_simple.py --docker --env 3DBall --episodes 1070 71=============================================================================72 73The first run will download Unity environment binaries (~500MB).74Subsequent runs use cached binaries from ~/.mlagents-cache/75"""76 77import argparse78import random79import sys80import time81from pathlib import Path82from typing import Optional83 84# Add paths for imports85sys.path.insert(0, str(Path(__file__).parent.parent / "src"))86sys.path.insert(0, str(Path(__file__).parent.parent))87 88from envs.unity_env.client import UnityEnv89from envs.unity_env.models import UnityAction90 91 92def run_pushblock_episode(93 client: UnityEnv,94 max_steps: int = 1000,95 verbose: bool = True,96) -> dict:97 """98 Run a single episode of PushBlock with random actions.99 100 Args:101 client: Connected UnityEnv client102 max_steps: Maximum steps per episode103 verbose: Print progress information104 105 Returns:106 Dictionary with episode statistics107 """108 # Reset to PushBlock environment109 result = client.reset(env_id="PushBlock")110 111 if verbose:112 print(f"Environment: PushBlock")113 print(f"Behavior: {result.observation.behavior_name}")114 print(f"Vector obs dims: {len(result.observation.vector_observations)}")115 action_spec = result.observation.action_spec_info116 print(f"Action spec: {action_spec}")117 print()118 119 episode_reward = 0.0120 step_count = 0121 122 while not result.done and step_count < max_steps:123 # PushBlock has 7 discrete actions:124 # 0=noop, 1=forward, 2=backward, 3=rotate_left,125 # 4=rotate_right, 5=strafe_left, 6=strafe_right126 action_idx = random.randint(0, 6)127 action = UnityAction(discrete_actions=[action_idx])128 129 result = client.step(action)130 episode_reward += result.reward or 0.0131 step_count += 1132 133 if verbose and step_count % 100 == 0:134 print(f" Step {step_count}: cumulative reward = {episode_reward:.2f}")135 136 return {137 "steps": step_count,138 "reward": episode_reward,139 "done": result.done,140 }141 142 143def run_3dball_episode(144 client: UnityEnv,145 max_steps: int = 500,146 verbose: bool = True,147) -> dict:148 """149 Run a single episode of 3DBall with random actions.150 151 Args:152 client: Connected UnityEnv client153 max_steps: Maximum steps per episode154 verbose: Print progress information155 156 Returns:157 Dictionary with episode statistics158 """159 # Reset to 3DBall environment160 result = client.reset(env_id="3DBall")161 162 if verbose:163 print(f"Environment: 3DBall")164 print(f"Behavior: {result.observation.behavior_name}")165 print(f"Vector obs dims: {len(result.observation.vector_observations)}")166 action_spec = result.observation.action_spec_info167 print(f"Action spec: {action_spec}")168 print()169 170 episode_reward = 0.0171 step_count = 0172 173 while not result.done and step_count < max_steps:174 # 3DBall has 2 continuous actions for X and Z rotation175 action = UnityAction(176 continuous_actions=[177 random.uniform(-1, 1), # X rotation178 random.uniform(-1, 1), # Z rotation179 ]180 )181 182 result = client.step(action)183 episode_reward += result.reward or 0.0184 step_count += 1185 186 if verbose and step_count % 100 == 0:187 print(f" Step {step_count}: cumulative reward = {episode_reward:.2f}")188 189 return {190 "steps": step_count,191 "reward": episode_reward,192 "done": result.done,193 }194 195 196def run_episodes(197 client: UnityEnv,198 env_name: str,199 episodes: int,200 max_steps: int,201 verbose: bool,202) -> list:203 """Run multiple episodes and collect results."""204 all_results = []205 206 for episode in range(episodes):207 print(f"\n--- Episode {episode + 1}/{episodes} ---")208 209 if env_name == "PushBlock":210 result = run_pushblock_episode(211 client,212 max_steps=max_steps,213 verbose=verbose,214 )215 elif env_name == "3DBall":216 result = run_3dball_episode(217 client,218 max_steps=max_steps,219 verbose=verbose,220 )221 else: # both222 if episode % 2 == 0:223 result = run_pushblock_episode(224 client,225 max_steps=max_steps,226 verbose=verbose,227 )228 else:229 result = run_3dball_episode(230 client,231 max_steps=max_steps,232 verbose=verbose,233 )234 235 all_results.append(result)236 print(237 f"Episode {episode + 1}: {result['steps']} steps, "238 f"reward: {result['reward']:.2f}"239 )240 241 return all_results242 243 244def print_summary(all_results: list) -> None:245 """Print summary statistics."""246 print("\n" + "=" * 60)247 print("Summary")248 print("=" * 60)249 total_steps = sum(r["steps"] for r in all_results)250 avg_reward = sum(r["reward"] for r in all_results) / len(all_results)251 max_reward = max(r["reward"] for r in all_results)252 min_reward = min(r["reward"] for r in all_results)253 print(f"Total episodes: {len(all_results)}")254 print(f"Total steps: {total_steps}")255 print(f"Average reward: {avg_reward:.2f}")256 print(f"Max reward: {max_reward:.2f}")257 print(f"Min reward: {min_reward:.2f}")258 print("=" * 60)259 260 261def run_with_server(args) -> None:262 """Run using a connection to an existing server."""263 print("=" * 60)264 print("Unity ML-Agents Environment - Server Mode")265 print("=" * 60)266 print(f"\nConnecting to: {args.url}")267 print(f"Environment: {args.env}")268 print(f"Episodes: {args.episodes}")269 print(f"Max steps: {args.max_steps}")270 print()271 272 # Connect to the environment server273 with UnityEnv(base_url=args.url) as client:274 all_results = run_episodes(275 client,276 env_name=args.env,277 episodes=args.episodes,278 max_steps=args.max_steps,279 verbose=not args.quiet,280 )281 print_summary(all_results)282 283 284def run_with_docker(args) -> None:285 """Run using Docker (automatically starts container)."""286 print("=" * 60)287 print("Unity ML-Agents Environment - Docker Mode")288 print("=" * 60)289 print(f"\nDocker image: {args.docker_image}")290 print(f"Environment: {args.env}")291 print(f"Episodes: {args.episodes}")292 print(f"Max steps: {args.max_steps}")293 print(f"Window size: {args.width}x{args.height}")294 print(f"Graphics: {'Disabled (headless)' if args.no_graphics else 'Enabled'}")295 print()296 297 # Build environment variables for Docker298 env_vars = {299 "UNITY_NO_GRAPHICS": "1" if args.no_graphics else "0",300 "UNITY_WIDTH": str(args.width),301 "UNITY_HEIGHT": str(args.height),302 "UNITY_TIME_SCALE": str(args.time_scale),303 "UNITY_QUALITY_LEVEL": str(args.quality_level),304 }305 306 print("Starting Docker container...")307 print(f" Environment variables: {env_vars}")308 print()309 310 try:311 # Use from_docker_image to automatically start and connect312 client = UnityEnv.from_docker_image(313 args.docker_image,314 environment=env_vars,315 )316 317 try:318 all_results = run_episodes(319 client,320 env_name=args.env,321 episodes=args.episodes,322 max_steps=args.max_steps,323 verbose=not args.quiet,324 )325 print_summary(all_results)326 finally:327 print("\nClosing Docker container...")328 client.close()329 330 except Exception as e:331 print(f"\nError running with Docker: {e}")332 print("\nTroubleshooting:")333 print(" 1. Ensure Docker is running")334 print(" 2. Build the image first:")335 print(f" docker build -f server/Dockerfile -t {args.docker_image} .")336 print(" 3. Or use server mode instead:")337 print(" python examples/unity_simple.py --url http://localhost:8000")338 sys.exit(1)339 340 341def run_direct(args) -> None:342 """343 Run Unity environment in direct mode (local server started automatically).344 345 This mode starts an embedded local server and connects to it, providing346 the convenience of direct execution while maintaining client-server separation.347 Useful for quick testing and debugging. For production, use server or Docker mode.348 """349 print("=" * 60)350 print("Unity ML-Agents Environment - Direct Mode")351 print("=" * 60)352 print(f"\nEnvironment: {args.env}")353 print(f"Episodes: {args.episodes}")354 print(f"Max steps: {args.max_steps}")355 print(f"Window size: {args.width}x{args.height}")356 print(f"Graphics: {'Disabled (headless)' if args.no_graphics else 'Enabled'}")357 print(f"Time scale: {args.time_scale}x")358 print()359 360 print("Starting local Unity server...")361 print("(First run will download binaries - this may take a few minutes)")362 print()363 364 # Use from_direct() to start an embedded server and get a client365 client = UnityEnv.from_direct(366 env_id=args.env if args.env != "both" else "PushBlock",367 no_graphics=args.no_graphics,368 time_scale=args.time_scale,369 width=args.width,370 height=args.height,371 quality_level=args.quality_level,372 )373 374 try:375 all_results = []376 377 for episode in range(args.episodes):378 print(f"\n--- Episode {episode + 1}/{args.episodes} ---")379 380 # Determine which environment to use381 if args.env == "both":382 current_env = "PushBlock" if episode % 2 == 0 else "3DBall"383 else:384 current_env = args.env385 386 # Reset environment387 result = client.reset(env_id=current_env)388 389 if not args.quiet:390 print(f"Environment: {current_env}")391 print(f"Behavior: {result.observation.behavior_name}")392 print(f"Vector obs dims: {len(result.observation.vector_observations)}")393 print(f"Action spec: {result.observation.action_spec_info}")394 print()395 396 episode_reward = 0.0397 step_count = 0398 399 while not result.done and step_count < args.max_steps:400 # Generate action based on environment type401 if current_env == "3DBall":402 action = UnityAction(403 continuous_actions=[404 random.uniform(-1, 1),405 random.uniform(-1, 1),406 ]407 )408 else:409 action = UnityAction(discrete_actions=[random.randint(0, 6)])410 411 result = client.step(action)412 episode_reward += result.reward or 0.0413 step_count += 1414 415 if not args.quiet and step_count % 100 == 0:416 print(417 f" Step {step_count}: cumulative reward = {episode_reward:.2f}"418 )419 420 episode_result = {421 "steps": step_count,422 "reward": episode_reward,423 "done": result.done,424 }425 all_results.append(episode_result)426 print(427 f"Episode {episode + 1}: {episode_result['steps']} steps, "428 f"reward: {episode_result['reward']:.2f}"429 )430 431 print_summary(all_results)432 433 finally:434 print("\nClosing Unity environment...")435 client.close()436 437 438def main():439 parser = argparse.ArgumentParser(440 description="Run Unity ML-Agents environment examples",441 formatter_class=argparse.RawDescriptionHelpFormatter,442 epilog="""443Examples:444 # Connect to running server (default)445 %(prog)s --url http://localhost:8000446 447 # Run via Docker448 %(prog)s --docker449 450 # Run directly without server (for testing)451 %(prog)s --direct452 453 # With graphics window (800x600 default)454 %(prog)s --direct --width 1280 --height 720455 456 # Headless mode (faster training)457 %(prog)s --direct --no-graphics --time-scale 20458 459 # Run 3DBall environment460 %(prog)s --direct --env 3DBall --episodes 5461 """,462 )463 464 # Mode selection465 mode_group = parser.add_mutually_exclusive_group()466 mode_group.add_argument(467 "--docker",468 action="store_true",469 help="Run via Docker (automatically starts container)",470 )471 mode_group.add_argument(472 "--direct",473 action="store_true",474 help="Run Unity environment directly without server",475 )476 477 # Connection settings478 parser.add_argument(479 "--url",480 default="http://localhost:8000",481 help="Base URL of the Unity environment server (default: http://localhost:8000)",482 )483 parser.add_argument(484 "--docker-image",485 default="unity-env:latest",486 help="Docker image to use (default: unity-env:latest)",487 )488 489 # Environment settings490 parser.add_argument(491 "--env",492 choices=["PushBlock", "3DBall", "both"],493 default="PushBlock",494 help="Which environment to run (default: PushBlock)",495 )496 parser.add_argument(497 "--episodes",498 type=int,499 default=3,500 help="Number of episodes to run (default: 3)",501 )502 parser.add_argument(503 "--max-steps",504 type=int,505 default=500,506 help="Maximum steps per episode (default: 500)",507 )508 509 # Graphics settings510 parser.add_argument(511 "--width",512 type=int,513 default=1280,514 help="Window width in pixels (default: 800)",515 )516 parser.add_argument(517 "--height",518 type=int,519 default=720,520 help="Window height in pixels (default: 600)",521 )522 parser.add_argument(523 "--no-graphics",524 action="store_true",525 help="Run in headless mode without graphics (faster training)",526 )527 parser.add_argument(528 "--time-scale",529 type=float,530 default=1.0,531 help="Simulation speed multiplier (default: 1.0, use 20.0 for fast training)",532 )533 parser.add_argument(534 "--quality-level",535 type=int,536 default=5,537 choices=[0, 1, 2, 3, 4, 5],538 help="Graphics quality level 0-5 (default: 5)",539 )540 541 # Output settings542 parser.add_argument(543 "--quiet",544 action="store_true",545 help="Reduce output verbosity",546 )547 548 args = parser.parse_args()549 550 # Run in appropriate mode551 if args.docker:552 run_with_docker(args)553 elif args.direct:554 run_direct(args)555 else:556 run_with_server(args)557 558 559if __name__ == "__main__":560 main()561