Spaces:
Sleeping
Sleeping
| """Watch a random agent drop bricks - the M1 demo. | |
| uv run python -m baselines.random_agent --render human --episodes 3 | |
| Headless? Use --gif out.gif instead (or SDL_VIDEODRIVER=dummy). | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import gymnasium as gym | |
| import atrium_sim # noqa: F401 (registers the env) | |
| from baselines.policy import RandomPolicy | |
| def main() -> None: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--render", choices=["human", "rgb_array", "none"], default="human") | |
| parser.add_argument("--episodes", type=int, default=3) | |
| parser.add_argument("--seed", type=int, default=1) | |
| parser.add_argument("--suite", default="train") | |
| parser.add_argument("--gif", default=None, help="record the first episode to this GIF path") | |
| args = parser.parse_args() | |
| render_mode = None if args.render == "none" else args.render | |
| env = gym.make("atrium_sim/BrickLayer-v0", render_mode=render_mode) | |
| env.unwrapped.env_cfg = type(env.unwrapped.env_cfg)(suite=args.suite) | |
| policy = RandomPolicy(env.action_space, args.seed) | |
| for ep in range(args.episodes): | |
| if args.gif and ep == 0: | |
| from atrium_sim.render.recorder import record_episode | |
| metrics = record_episode(env, policy, args.gif, seed=args.seed) | |
| else: | |
| obs, info = env.reset(seed=args.seed + ep) | |
| done = False | |
| while not done: | |
| obs, reward, terminated, truncated, info = env.step(policy.act(obs)) | |
| done = terminated or truncated | |
| metrics = info["metrics"] | |
| print( | |
| f"episode {ep}: return={metrics['episode_return']:+.2f} " | |
| f"in-tol={metrics['frac_in_tol']:.0%} filled={metrics['frac_filled']:.0%} " | |
| f"waste={metrics['waste_count']:.0f} placements={metrics['placements']:.0f}" | |
| ) | |
| env.close() | |
| if __name__ == "__main__": | |
| main() | |