| import time, numpy as np, gymnasium as gym, ogbench, imageio |
| print("=== visual singletask env ===") |
| env = gym.make("visual-cube-single-singletask-task1-v0") |
| print("obs_space:", env.observation_space) |
| print("act_space:", env.action_space) |
| ob, info = env.reset(seed=0) |
| ob = np.asarray(ob) |
| print("obs shape", ob.shape, ob.dtype, "min", ob.min(), "max", ob.max()) |
| imageio.imwrite("sample_env_frame.png", ob if ob.ndim==3 else ob[...,:3]) |
| t0=time.time(); n=200 |
| for i in range(n): |
| ob,r,te,tr,info = env.step(env.action_space.sample()) |
| if te or tr: env.reset(seed=i) |
| dt=time.time()-t0 |
| print(f"render+step speed: {n/dt:.1f} steps/s -> 200-step episode = {200/(n/dt):.1f}s") |
| print("=== goal-conditioned visual env ===") |
| g = gym.make("visual-cube-single-v0") |
| ob2, info2 = g.reset(seed=0, options=dict(task_id=1, render_goal=True)) |
| print("info keys:", [k for k in info2.keys() if "goal" in k]) |
| gi = info2.get("goal") |
| print("goal shape", None if gi is None else np.asarray(gi).shape) |
| if gi is not None: imageio.imwrite("sample_goal_frame.png", np.asarray(gi)) |
|
|
