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))