(task_name, *, steps, seed, out_dir, side_by_side)
| 86 | |
| 87 | |
| 88 | def render_task(task_name, *, steps, seed, out_dir, side_by_side): |
| 89 | import gymnasium as gym |
| 90 | import simpler_env |
| 91 | |
| 92 | from roboverse_pack.tasks.simpler_env import register_simpler_env_passthrough |
| 93 | |
| 94 | register_simpler_env_passthrough() |
| 95 | |
| 96 | wrapped = gym.make(f"SimplerEnv/{task_name}") |
| 97 | rng = np.random.RandomState(seed) |
| 98 | space = wrapped.action_space |
| 99 | actions = [rng.uniform(space.low, space.high).astype(space.dtype) for _ in range(steps)] |
| 100 | |
| 101 | fw = _rollout_frames(wrapped, actions, seed) |
| 102 | wrapped.close() |
| 103 | |
| 104 | os.makedirs(out_dir, exist_ok=True) |
| 105 | out_path = os.path.join(out_dir, f"{task_name}.mp4") |
| 106 | |
| 107 | if side_by_side: |
| 108 | native = simpler_env.make(task_name) |
| 109 | fn = _rollout_frames(native, actions, seed) |
| 110 | native.close() |
| 111 | n = min(len(fn), len(fw)) |
| 112 | combo = [] |
| 113 | for i in range(n): |
| 114 | a = np.asarray(fn[i]) |
| 115 | b = np.asarray(fw[i]) |
| 116 | if a.shape == b.shape: |
| 117 | combo.append(np.concatenate([a, b], axis=1)) |
| 118 | if combo: |
| 119 | ok = _save_mp4(combo, out_path) |
| 120 | return out_path if ok else None |
| 121 | |
| 122 | ok = _save_mp4(fw, out_path) |
| 123 | return out_path if ok else None |
| 124 | |
| 125 | |
| 126 | def main() -> None: |
no test coverage detected