JackM323
c5ca79a354
Training environment to make a walk model for the hexapod generated code that will be checked
79 lines
2.6 KiB
Python
79 lines
2.6 KiB
Python
import argparse
|
|
import numpy as np
|
|
|
|
from .env import JackBotEnv
|
|
|
|
|
|
def evaluate(
|
|
model_path: str,
|
|
episodes: int = 5,
|
|
use_gui: bool = False,
|
|
num_robots: int = 1,
|
|
robot_spacing: float = 0.5,
|
|
start_pose: str = "init_deg",
|
|
):
|
|
try:
|
|
from stable_baselines3 import PPO
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"stable-baselines3 is required for evaluation. Install with: pip install stable-baselines3"
|
|
) from exc
|
|
|
|
env = JackBotEnv(
|
|
use_gui=use_gui,
|
|
random_command=False,
|
|
num_robots=num_robots,
|
|
robot_spacing=robot_spacing,
|
|
start_pose=start_pose,
|
|
)
|
|
model = PPO.load(model_path)
|
|
|
|
for episode in range(episodes):
|
|
reset_res = env.reset()
|
|
# handle Gym / Gymnasium compatibility: reset may return (obs, info)
|
|
if isinstance(reset_res, tuple) and len(reset_res) == 2:
|
|
obs, _ = reset_res
|
|
else:
|
|
obs = reset_res
|
|
|
|
done = False
|
|
episode_reward = 0.0
|
|
|
|
while not done:
|
|
# pass only the observation to the policy
|
|
action, _ = model.predict(obs, deterministic=True)
|
|
|
|
step_res = env.step(action)
|
|
# Gymnasium-style: (obs, reward, terminated, truncated, info)
|
|
if isinstance(step_res, tuple) and len(step_res) == 5:
|
|
obs, reward, terminated, truncated, info = step_res
|
|
done = bool(terminated or truncated)
|
|
else:
|
|
# legacy Gym: (obs, reward, done, info)
|
|
obs, reward, done, info = step_res
|
|
|
|
episode_reward += float(reward)
|
|
|
|
print(f"Episode {episode + 1}: reward={episode_reward:.2f}")
|
|
|
|
env.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Evaluate a trained JackBot policy.")
|
|
parser.add_argument("--model-path", type=str, required=True)
|
|
parser.add_argument("--episodes", type=int, default=5)
|
|
parser.add_argument("--gui", action="store_true")
|
|
parser.add_argument("--num-robots", type=int, default=1, help="Number of robots in the environment")
|
|
parser.add_argument("--robot-spacing", type=float, default=0.5, help="Spacing between robots in meters")
|
|
parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
|
|
args = parser.parse_args()
|
|
evaluate(
|
|
args.model_path,
|
|
episodes=args.episodes,
|
|
use_gui=args.gui,
|
|
num_robots=args.num_robots,
|
|
robot_spacing=args.robot_spacing,
|
|
start_pose=args.start_pose,
|
|
)
|