Branch data Line data Source code
1 : : #!/usr/bin/env python3
2 : :
3 : : # Copyright 2026 Garena Online Private Limited
4 : : #
5 : : # Licensed under the Apache License, Version 2.0 (the "License");
6 : : # you may not use this file except in compliance with the License.
7 : : # You may obtain a copy of the License at
8 : : #
9 : : # http://www.apache.org/licenses/LICENSE-2.0
10 : : #
11 : : # Unless required by applicable law or agreed to in writing, software
12 : : # distributed under the License is distributed on an "AS IS" BASIS,
13 : : # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 : : # See the License for the specific language governing permissions and
15 : : # limitations under the License.
16 : : """Generate render comparison and sample images for docs."""
17 : :
18 : 1 : from __future__ import annotations
19 : :
20 : 1 : import argparse
21 : 1 : import math
22 : 1 : import os
23 : 1 : from collections.abc import Callable
24 : 1 : from dataclasses import dataclass
25 : 1 : from pathlib import Path
26 : 1 : from typing import Any
27 : :
28 : 1 : import numpy as np
29 : 1 : from PIL import Image, ImageDraw, ImageFont
30 : :
31 : 1 : from envpool.mujoco.render_oracle import configure_linux_mujoco_gl
32 : 1 : from envpool.python.glfw_context import preload_windows_gl_dlls
33 : :
34 : 1 : preload_windows_gl_dlls(strict=True)
35 : 1 : configure_linux_mujoco_gl()
36 : :
37 : :
38 : 1 : @dataclass(frozen=True)
39 : 1 : class RenderCompareConfig:
40 : : """Shared image generation and comparison settings."""
41 : :
42 : 1 : family: str
43 : 1 : tile_width: int
44 : 1 : tile_height: int
45 : 1 : source_width: int
46 : 1 : source_height: int
47 : 1 : columns: int
48 : 1 : seed: int
49 : 1 : camera_id: int
50 : 1 : max_mean_abs_diff: float
51 : 1 : max_mismatch_ratio: float
52 : 1 : max_ignored_abs_diff: int
53 : 1 : require_bitwise: bool
54 : 1 : flip_vertical: bool
55 : :
56 : :
57 : 1 : RenderPairFn = Callable[
58 : : [str, RenderCompareConfig], tuple[np.ndarray, np.ndarray | None]
59 : : ]
60 : :
61 : :
62 : 1 : @dataclass(frozen=True)
63 : 1 : class RenderItem:
64 : : """One family-specific render key and its display label."""
65 : :
66 : 1 : key: str
67 : 1 : label: str
68 : :
69 : :
70 : 1 : @dataclass(frozen=True)
71 : 1 : class RenderFamily:
72 : : """Family-specific render integration."""
73 : :
74 : 1 : items: tuple[RenderItem, ...]
75 : 1 : default_output: Path
76 : 1 : render_pair: RenderPairFn
77 : 1 : left_title: str = "EnvPool"
78 : 1 : right_title: str | None = "Official"
79 : 1 : default_flip_vertical: bool = False
80 : 1 : compare_frames: bool = True
81 : :
82 : :
83 : 1 : _PAIR_GAP = 4
84 : 1 : _CELL_GAP = 12
85 : 1 : _HEADER_HEIGHT = 28
86 : 1 : _MARGIN = 16
87 : :
88 : :
89 : 1 : def _make_craftax_family() -> RenderFamily:
90 : : """Build the Craftax comparison adapter using actual gameplay frames."""
91 : 0 : import envpool.craftax.registration # noqa: F401
92 : 0 : from envpool.craftax.oracle import (
93 : : encode,
94 : : jax,
95 : : make_oracle,
96 : : native,
97 : : renderer,
98 : : )
99 : 0 : from envpool.registration import make_gymnasium
100 : :
101 : 0 : scenes = {
102 : : "Classic": (True, 0, 96),
103 : : "Craftax": (False, 0, 96),
104 : : "Cave": (False, 2, 48),
105 : : }
106 : :
107 : 0 : def render_pair(
108 : : scene: str, cfg: RenderCompareConfig
109 : : ) -> tuple[np.ndarray, np.ndarray]:
110 : 0 : classic, level, length = scenes[scene]
111 : 0 : task = (
112 : : "Craftax-"
113 : : + ("Classic-" if classic else "")
114 : : + "Symbolic-AutoReset-v1"
115 : : )
116 : 0 : oracle, params = make_oracle(task)
117 : 0 : key = jax.random.PRNGKey(cfg.seed)
118 : 0 : _, state = oracle.reset(key, params)
119 : 0 : kwargs = {}
120 [ # # ]: 0 : if level:
121 : 0 : state = state.replace(
122 : : player_level=jax.numpy.asarray(level, dtype=jax.numpy.int32),
123 : : player_position=state.up_ladders[level],
124 : : )
125 : 0 : layout = native.Game(native.Params(classic)).get_state()
126 : 0 : kwargs["initial_state"] = encode(state, layout).tolist()
127 : 0 : pool = make_gymnasium(
128 : : task, num_envs=1, seed=cfg.seed, render_mode="rgb_array", **kwargs
129 : : )
130 : 0 : step = jax.jit(oracle.step)
131 : 0 : rng = np.random.default_rng(40)
132 : 0 : try:
133 : 0 : pool.reset()
134 [ # # ]: 0 : for _ in range(length):
135 : 0 : key, draw = jax.random.split(key)
136 : 0 : action = int(rng.integers(17 if classic else 43))
137 : 0 : _, state, _, _, _ = step(draw, state, action, params)
138 : 0 : pool.step(np.array([action], np.int32))
139 : 0 : actual = np.asarray(pool.render())[0]
140 : 0 : expected = np.asarray(renderer(classic)(state)).astype(np.uint8)
141 : 0 : np.testing.assert_array_equal(actual, expected)
142 : 0 : return actual, expected
143 : : finally:
144 : 0 : pool.close()
145 : :
146 : 0 : return RenderFamily(
147 : : items=tuple(
148 : : RenderItem(key=scene, label=f"{scene} / {length} actions")
149 : : for scene, (_, _, length) in scenes.items()
150 : : ),
151 : : default_output=Path(
152 : : "docs/_static/render_samples/craftax_official_compare.png"
153 : : ),
154 : : render_pair=render_pair,
155 : : right_title="Official Craftax 1.6.1",
156 : : )
157 : :
158 : :
159 : 1 : def _make_metaworld_family() -> RenderFamily:
160 : : """Build the MetaWorld render comparison adapter."""
161 : 0 : import mujoco
162 : 0 : from metaworld.env_dict import ALL_V3_ENVIRONMENTS
163 : :
164 : 0 : import envpool.mujoco.metaworld.registration as metaworld_registration
165 : 0 : from envpool.registration import make_gymnasium
166 : :
167 : 0 : def make_oracle(task_name: str, cfg: RenderCompareConfig) -> Any:
168 : 0 : env = ALL_V3_ENVIRONMENTS[task_name](
169 : : render_mode="rgb_array",
170 : : camera_id=cfg.camera_id,
171 : : )
172 : 0 : env._set_task_called = True
173 : 0 : env._partially_observable = True
174 : 0 : env.mujoco_renderer.width = cfg.source_width
175 : 0 : env.mujoco_renderer.height = cfg.source_height
176 : 0 : return env
177 : :
178 : 0 : def sync_reset_state(oracle: Any, info: dict[str, Any]) -> None:
179 : 0 : required_keys = (
180 : : "rand_vec0",
181 : : "qpos0",
182 : : "qvel0",
183 : : "mocap_pos0",
184 : : "mocap_quat0",
185 : : "qacc0",
186 : : "qacc_warmstart0",
187 : : "init_tcp0",
188 : : "init_left_pad0",
189 : : "init_right_pad0",
190 : : )
191 : 0 : missing = [key for key in required_keys if key not in info]
192 [ # # ]: 0 : if missing:
193 : 0 : raise RuntimeError(
194 : : "MetaWorld render comparison must be generated with "
195 : : "`bazel run --config=debug //scripts:render_compare -- "
196 : : "--family=metaworld` so EnvPool exposes reset-sync info. "
197 : : f"Missing keys: {missing}"
198 : : )
199 : :
200 : 0 : rand_vec = np.asarray(info["rand_vec0"][0], dtype=np.float64)
201 : 0 : random_dim = int(oracle._random_reset_space.low.size)
202 : 0 : oracle._freeze_rand_vec = True
203 : 0 : oracle._last_rand_vec = rand_vec[:random_dim].copy()
204 : 0 : oracle.reset()
205 : :
206 : 0 : qpos = np.asarray(info["qpos0"][0], dtype=np.float64)[
207 : : : oracle.data.qpos.size
208 : : ]
209 : 0 : qvel = np.asarray(info["qvel0"][0], dtype=np.float64)[
210 : : : oracle.data.qvel.size
211 : : ]
212 : 0 : oracle.set_state(qpos, qvel)
213 : 0 : oracle.data.mocap_pos[0] = np.asarray(
214 : : info["mocap_pos0"][0], dtype=np.float64
215 : : )
216 : 0 : oracle.data.mocap_quat[0] = np.asarray(
217 : : info["mocap_quat0"][0], dtype=np.float64
218 : : )
219 : 0 : oracle.data.qacc[:] = np.asarray(info["qacc0"][0], dtype=np.float64)[
220 : : : oracle.data.qacc.size
221 : : ]
222 : 0 : oracle.data.qacc_warmstart[:] = np.asarray(
223 : : info["qacc_warmstart0"][0], dtype=np.float64
224 : : )[: oracle.data.qacc_warmstart.size]
225 : 0 : mujoco.mj_forward(oracle.model, oracle.data)
226 : 0 : oracle.init_tcp = np.asarray(
227 : : info["init_tcp0"][0], dtype=np.float64
228 : : ).copy()
229 : 0 : oracle.init_left_pad = np.asarray(
230 : : info["init_left_pad0"][0], dtype=np.float64
231 : : ).copy()
232 : 0 : oracle.init_right_pad = np.asarray(
233 : : info["init_right_pad0"][0], dtype=np.float64
234 : : ).copy()
235 [ # # ]: 0 : if hasattr(oracle, "_handle_init_pos"):
236 : 0 : oracle._handle_init_pos = oracle._get_pos_objects().copy()
237 : :
238 : 0 : curr_obs = oracle._get_curr_obs_combined_no_goal()
239 : 0 : oracle._prev_obs = curr_obs.copy()
240 : 0 : obs = oracle._get_obs().astype(np.float64)
241 : 0 : oracle._last_stable_obs = obs.copy()
242 : :
243 : 0 : def render_pair(
244 : : task_name: str,
245 : : cfg: RenderCompareConfig,
246 : : ) -> tuple[np.ndarray, np.ndarray]:
247 : 0 : task_id = metaworld_registration.metaworld_task_id(task_name)
248 : 0 : env = make_gymnasium(
249 : : task_id,
250 : : num_envs=1,
251 : : seed=cfg.seed,
252 : : render_mode="rgb_array",
253 : : render_width=cfg.source_width,
254 : : render_height=cfg.source_height,
255 : : render_camera_id=cfg.camera_id,
256 : : )
257 : 0 : oracle = make_oracle(task_name, cfg)
258 : 0 : try:
259 : 0 : _, info = env.reset()
260 : 0 : sync_reset_state(oracle, info)
261 : 0 : envpool_frame = env.render()
262 [ # # ]: 0 : if envpool_frame is None:
263 : 0 : raise RuntimeError(f"{task_id} returned no EnvPool frame")
264 : 0 : return envpool_frame[0], oracle.render()
265 : : finally:
266 : 0 : env.close()
267 : 0 : oracle.close()
268 : :
269 : 0 : return RenderFamily(
270 : : items=tuple(
271 : : RenderItem(
272 : : key=task_name,
273 : : label=metaworld_registration.metaworld_public_task_name(
274 : : task_name
275 : : ),
276 : : )
277 : : for task_name in metaworld_registration.metaworld_v3_envs
278 : : ),
279 : : default_output=Path(
280 : : "docs/_static/render_samples/metaworld_official_compare.png"
281 : : ),
282 : : render_pair=render_pair,
283 : : default_flip_vertical=True,
284 : : )
285 : :
286 : :
287 : 1 : def _make_jumanji_family() -> RenderFamily:
288 : : """Use the same reset synchronization as independent alignment tests."""
289 : 0 : import envpool.jumanji.registration as jumanji_registration
290 : 0 : from envpool.jumanji import oracle
291 : 0 : from envpool.registration import make_gymnasium
292 : :
293 : 0 : def render_pair(
294 : : task_id: str, cfg: RenderCompareConfig
295 : : ) -> tuple[np.ndarray, np.ndarray]:
296 : 0 : official_env = oracle._make_official_env(task_id)
297 : 0 : env = None
298 : 0 : try:
299 : 0 : state, timestep = official_env.reset(
300 : : oracle.jax.random.PRNGKey(cfg.seed)
301 : : )
302 : 0 : kwargs = oracle._envpool_kwargs_from_official_state(
303 : : task_id, state, timestep.observation
304 : : )
305 [ # # ]: 0 : if task_id == "Sokoban-v0":
306 : 0 : kwargs.update(
307 : : base_path="/tmp/envpool-missing-boxoban",
308 : : sokoban_level_index=0,
309 : : )
310 : 0 : env = make_gymnasium(
311 : : task_id,
312 : : num_envs=1,
313 : : seed=cfg.seed,
314 : : render_mode="rgb_array",
315 : : render_width=cfg.source_width,
316 : : render_height=cfg.source_height,
317 : : **kwargs,
318 : : )
319 : 0 : env.reset()
320 : 0 : frames = env.render()
321 : 0 : assert frames is not None
322 : 0 : return frames[0], oracle._official_render(
323 : : official_env, state, cfg.source_width, cfg.source_height
324 : : )
325 : : finally:
326 [ # # ]: 0 : if env is not None:
327 : 0 : env.close()
328 : 0 : official_env.close()
329 : :
330 : 0 : return RenderFamily(
331 : : items=tuple(
332 : : RenderItem(key=task_id, label=task_id)
333 : : for task_id in jumanji_registration.jumanji_env_ids
334 : : ),
335 : : default_output=Path(
336 : : "docs/_static/render_samples/jumanji_official_compare.png"
337 : : ),
338 : : render_pair=render_pair,
339 : : left_title="EnvPool",
340 : : right_title="Official",
341 : : # Native task graphics intentionally use a different layout and style.
342 : : compare_frames=False,
343 : : )
344 : :
345 : :
346 : 1 : def _make_mujoco_playground_family() -> RenderFamily:
347 : : """Build the MuJoCo Playground official render comparison adapter."""
348 : 1 : from jax import config as jax_config
349 : :
350 : 1 : jax_config.update("jax_enable_x64", True)
351 : :
352 : 1 : import mujoco
353 : :
354 : 1 : import envpool.mujoco.playground.registration as playground_registration
355 : 1 : from envpool.mujoco.playground.oracle import (
356 : : FrameRenderer,
357 : : make_oracle,
358 : : native_config,
359 : : )
360 : 1 : from envpool.mujoco.render_oracle import configure_macos_mujoco_renderer
361 : 1 : from envpool.registration import make_gymnasium
362 : :
363 : 1 : configure_macos_mujoco_renderer()
364 : :
365 : 1 : def make_render_data(model: Any, info: dict[str, Any]) -> Any:
366 : 1 : data = mujoco.MjData(model)
367 : 1 : data.qpos[:] = info["qpos"][0, : model.nq]
368 : 1 : data.qvel[:] = info["qvel"][0, : model.nv]
369 [ + - ]: 1 : if "ctrl" in info:
370 : 1 : data.ctrl[:] = info["ctrl"][0, : model.nu]
371 [ + + ]: 1 : if "mocap_pos" in info:
372 : 1 : data.mocap_pos[:] = info["mocap_pos"][
373 : : 0, : model.nmocap * 3
374 : : ].reshape(model.nmocap, 3)
375 [ + + ]: 1 : if "mocap_quat" in info:
376 : 1 : data.mocap_quat[:] = info["mocap_quat"][
377 : : 0, : model.nmocap * 4
378 : : ].reshape(model.nmocap, 4)
379 [ + - ]: 1 : if "qacc_warmstart" in info:
380 : 1 : data.qacc_warmstart[:] = info["qacc_warmstart"][0, : model.nv]
381 : 1 : mujoco.mj_forward(model, data)
382 [ + - ]: 1 : if "qacc" in info:
383 : 1 : data.qacc[:] = info["qacc"][0, : model.nv]
384 [ + - ]: 1 : if "qacc_warmstart" in info:
385 : 1 : data.qacc_warmstart[:] = info["qacc_warmstart"][0, : model.nv]
386 [ + - ]: 1 : if "sensordata" in info:
387 : 1 : data.sensordata[:] = info["sensordata"][0, : data.sensordata.size]
388 : 1 : return data
389 : :
390 : 1 : def render_pair(
391 : : task_id: str,
392 : : cfg: RenderCompareConfig,
393 : : ) -> tuple[np.ndarray, np.ndarray]:
394 : 1 : oracle = make_oracle(task_id)
395 : 1 : env = make_gymnasium(
396 : : task_id,
397 : : num_envs=1,
398 : : seed=cfg.seed,
399 : : render_mode="rgb_array",
400 : : render_width=cfg.source_width,
401 : : render_height=cfg.source_height,
402 : : **native_config(task_id),
403 : : )
404 : 1 : try:
405 : 1 : _, info = env.reset()
406 : 1 : render_model = oracle.mj_model
407 : 1 : render_data = make_render_data(render_model, info)
408 : 1 : renderer = FrameRenderer(
409 : : oracle,
410 : : task_id,
411 : : width=cfg.source_width,
412 : : height=cfg.source_height,
413 : : )
414 : 1 : try:
415 : 1 : official_frame = renderer.render(render_data)
416 : : finally:
417 : 1 : renderer.close()
418 : 1 : envpool_frame = env.render(env_ids=[0])
419 [ - + ]: 1 : if envpool_frame is None:
420 : 0 : raise RuntimeError(f"{task_id} returned no EnvPool frame")
421 : 1 : return np.asarray(envpool_frame[0]), official_frame
422 : : finally:
423 : 1 : env.close()
424 : :
425 : 1 : return RenderFamily(
426 : : items=tuple(
427 : : RenderItem(
428 : : key=f"{task_name}-v1",
429 : : label=task_name
430 : : .replace("BerkeleyHumanoid", "Berkeley ")
431 : : .replace("G1", "G1 ")
432 : : .replace("Go1", "Go1 ")
433 : : .replace("Op3", "OP3"),
434 : : )
435 : : for task_name in playground_registration.PLAYGROUND_ENVS
436 : : ),
437 : : default_output=Path(
438 : : "docs/_static/render_samples/mujoco_playground_official_compare.png"
439 : : ),
440 : : render_pair=render_pair,
441 : : )
442 : :
443 : :
444 : 1 : _FAMILY_BUILDERS: dict[str, Callable[[], RenderFamily]] = {
445 : : "craftax": _make_craftax_family,
446 : : "jumanji": _make_jumanji_family,
447 : : "metaworld": _make_metaworld_family,
448 : : "mujoco_playground": _make_mujoco_playground_family,
449 : : }
450 : :
451 : :
452 : 1 : def _make_display_image(
453 : : frame: np.ndarray, cfg: RenderCompareConfig
454 : : ) -> Image.Image:
455 : 1 : image = Image.fromarray(np.asarray(frame, dtype=np.uint8))
456 [ - + ]: 1 : if image.mode != "RGB":
457 : 0 : image = image.convert("RGB")
458 [ - + ]: 1 : if cfg.flip_vertical:
459 : 0 : image = image.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
460 : 1 : size = (cfg.tile_width, cfg.tile_height)
461 [ - + ]: 1 : if image.size == size:
462 : 0 : return image
463 : 1 : return image.resize(size, Image.Resampling.LANCZOS)
464 : :
465 : :
466 : 1 : def _render_mismatch_message(
467 : : label: str,
468 : : diff: np.ndarray,
469 : : mismatch_ratio: float,
470 : : max_ignored_abs_diff: int,
471 : : ) -> str:
472 : 0 : return (
473 : : f"{label} render mismatch: "
474 : : f"mean_abs_diff={float(diff.mean()):.6f}, "
475 : : f"max_abs_diff={int(diff.max())}, "
476 : : f"mismatch_ratio={mismatch_ratio:.6%}, "
477 : : f"ignored_abs_diff={max_ignored_abs_diff}"
478 : : )
479 : :
480 : :
481 : 1 : def _assert_frames_match(
482 : : label: str,
483 : : envpool_frame: np.ndarray,
484 : : official_frame: np.ndarray,
485 : : cfg: RenderCompareConfig,
486 : : ) -> None:
487 [ - + ]: 1 : if envpool_frame.shape != official_frame.shape:
488 : 0 : raise RuntimeError(
489 : : f"{label} render shape mismatch: "
490 : : f"{envpool_frame.shape} != {official_frame.shape}"
491 : : )
492 [ + - ]: 1 : if np.array_equal(envpool_frame, official_frame):
493 : 1 : return
494 : :
495 : 0 : diff = np.abs(
496 : : envpool_frame.astype(np.int16) - official_frame.astype(np.int16)
497 : : )
498 : 0 : mismatch_ratio = float(np.count_nonzero(diff)) / float(diff.size)
499 [ # # ]: 0 : if cfg.require_bitwise:
500 : 0 : raise RuntimeError(
501 : : _render_mismatch_message(label, diff, mismatch_ratio, 0)
502 : : )
503 : :
504 : 0 : mean_abs_diff = float(diff.mean())
505 : 0 : significant = diff > cfg.max_ignored_abs_diff
506 : 0 : mismatch_ratio = float(np.count_nonzero(significant)) / float(diff.size)
507 [ # # ]: 0 : if (
508 : : mean_abs_diff > cfg.max_mean_abs_diff
509 : : or mismatch_ratio > cfg.max_mismatch_ratio
510 : : ):
511 : 0 : raise RuntimeError(
512 : : _render_mismatch_message(
513 : : label, diff, mismatch_ratio, cfg.max_ignored_abs_diff
514 : : )
515 : : )
516 : :
517 : :
518 : 1 : def _draw_panel(
519 : : canvas: Image.Image,
520 : : family: RenderFamily,
521 : : label: str,
522 : : envpool_image: Image.Image,
523 : : official_image: Image.Image | None,
524 : : index: int,
525 : : cfg: RenderCompareConfig,
526 : : font: Any,
527 : : ) -> None:
528 : 1 : draw = ImageDraw.Draw(canvas)
529 : 1 : has_official = official_image is not None
530 : 1 : cell_width = cfg.tile_width * (2 if has_official else 1)
531 [ + - ]: 1 : if has_official:
532 : 1 : cell_width += _PAIR_GAP
533 : 1 : cell_height = _HEADER_HEIGHT + cfg.tile_height
534 : 1 : col = index % cfg.columns
535 : 1 : row = index // cfg.columns
536 : 1 : left = _MARGIN + col * (cell_width + _CELL_GAP)
537 : 1 : top = _MARGIN + row * (cell_height + _CELL_GAP)
538 : 1 : frame_top = top + _HEADER_HEIGHT
539 : :
540 : 1 : draw.text((left, top), label, fill=(30, 30, 30), font=font)
541 : 1 : draw.text((left, top + 14), family.left_title, fill=(80, 80, 80), font=font)
542 : 1 : canvas.paste(envpool_image, (left, frame_top))
543 : 1 : draw.rectangle(
544 : : [
545 : : left,
546 : : frame_top,
547 : : left + cfg.tile_width - 1,
548 : : frame_top + cfg.tile_height - 1,
549 : : ],
550 : : outline=(205, 205, 205),
551 : : )
552 [ + - ]: 1 : if has_official:
553 : 1 : official_left = left + cfg.tile_width + _PAIR_GAP
554 : 1 : assert family.right_title is not None
555 : 1 : draw.text(
556 : : (official_left, top + 14),
557 : : family.right_title,
558 : : fill=(80, 80, 80),
559 : : font=font,
560 : : )
561 : 1 : canvas.paste(official_image, (official_left, frame_top))
562 : 1 : draw.rectangle(
563 : : [
564 : : official_left,
565 : : frame_top,
566 : : official_left + cfg.tile_width - 1,
567 : : frame_top + cfg.tile_height - 1,
568 : : ],
569 : : outline=(205, 205, 205),
570 : : )
571 : :
572 : :
573 : 1 : def generate(
574 : : output: Path, family: RenderFamily, cfg: RenderCompareConfig
575 : : ) -> None:
576 : : """Generate and write one docs render image."""
577 : 1 : rows = math.ceil(len(family.items) / cfg.columns)
578 : 1 : has_official = family.right_title is not None
579 : 1 : cell_width = cfg.tile_width * (2 if has_official else 1)
580 [ + - ]: 1 : if has_official:
581 : 1 : cell_width += _PAIR_GAP
582 : 1 : cell_height = _HEADER_HEIGHT + cfg.tile_height
583 : 1 : width = _MARGIN * 2 + cfg.columns * cell_width
584 : 1 : width += (cfg.columns - 1) * _CELL_GAP
585 : 1 : height = _MARGIN * 2 + rows * cell_height + (rows - 1) * _CELL_GAP
586 : 1 : canvas = Image.new("RGB", (width, height), (255, 255, 255))
587 : 1 : font = ImageFont.load_default()
588 : :
589 [ + + ]: 1 : for index, item in enumerate(family.items):
590 : 1 : envpool_frame, official_frame = family.render_pair(item.key, cfg)
591 [ + - ]: 1 : if official_frame is not None and family.compare_frames:
592 : 1 : _assert_frames_match(item.label, envpool_frame, official_frame, cfg)
593 : 1 : envpool_image = _make_display_image(envpool_frame, cfg)
594 : 1 : official_image = (
595 : : None
596 : : if official_frame is None
597 : : else _make_display_image(official_frame, cfg)
598 : : )
599 : 1 : _draw_panel(
600 : : canvas,
601 : : family,
602 : : item.label,
603 : : envpool_image,
604 : : official_image,
605 : : index,
606 : : cfg,
607 : : font,
608 : : )
609 : :
610 : 1 : output.parent.mkdir(parents=True, exist_ok=True)
611 : 1 : canvas.save(output)
612 : :
613 : :
614 : 1 : def _parse_args() -> argparse.Namespace:
615 : 1 : parser = argparse.ArgumentParser()
616 : 1 : parser.add_argument(
617 : : "--family",
618 : : choices=sorted(_FAMILY_BUILDERS),
619 : : default="metaworld",
620 : : )
621 : 1 : parser.add_argument("--output", type=Path)
622 : 1 : parser.add_argument("--tile-width", type=int, default=96)
623 : 1 : parser.add_argument("--tile-height", type=int, default=72)
624 : 1 : parser.add_argument(
625 : : "--source-width",
626 : : type=int,
627 : : default=480,
628 : : help="Source render width before docs thumbnail downsampling.",
629 : : )
630 : 1 : parser.add_argument(
631 : : "--source-height",
632 : : type=int,
633 : : default=480,
634 : : help="Source render height before docs thumbnail downsampling.",
635 : : )
636 : 1 : parser.add_argument("--columns", type=int, default=5)
637 : 1 : parser.add_argument("--seed", type=int, default=7)
638 : 1 : parser.add_argument(
639 : : "--camera-id",
640 : : type=int,
641 : : default=1,
642 : : help="Renderer camera id for families that support camera selection.",
643 : : )
644 : 1 : parser.add_argument("--max-mean-abs-diff", type=float, default=0.25)
645 : 1 : parser.add_argument("--max-mismatch-ratio", type=float, default=0.005)
646 : 1 : parser.add_argument(
647 : : "--max-ignored-abs-diff",
648 : : type=int,
649 : : default=0,
650 : : help=(
651 : : "RGB channel deltas up to this value are ignored when computing "
652 : : "mismatch ratio. Mean absolute diff still uses raw deltas."
653 : : ),
654 : : )
655 : 1 : parser.add_argument("--require-bitwise", action="store_true")
656 : 1 : parser.add_argument(
657 : : "--flip-vertical",
658 : : action="store_true",
659 : : default=None,
660 : : help="Flip both EnvPool and oracle frames vertically in the output.",
661 : : )
662 : 1 : parser.add_argument(
663 : : "--no-flip-vertical",
664 : : action="store_false",
665 : : dest="flip_vertical",
666 : : help="Do not vertically flip output frames.",
667 : : )
668 : 1 : return parser.parse_args()
669 : :
670 : :
671 : 1 : def main() -> None:
672 : : """Parse command-line arguments and generate the comparison image."""
673 : 1 : args = _parse_args()
674 : 1 : family = _FAMILY_BUILDERS[args.family]()
675 : 1 : cfg = RenderCompareConfig(
676 : : family=args.family,
677 : : tile_width=args.tile_width,
678 : : tile_height=args.tile_height,
679 : : source_width=args.source_width,
680 : : source_height=args.source_height,
681 : : columns=args.columns,
682 : : seed=args.seed,
683 : : camera_id=args.camera_id,
684 : : max_mean_abs_diff=args.max_mean_abs_diff,
685 : : max_mismatch_ratio=args.max_mismatch_ratio,
686 : : max_ignored_abs_diff=args.max_ignored_abs_diff,
687 : : require_bitwise=args.require_bitwise,
688 : : flip_vertical=(
689 : : family.default_flip_vertical
690 : : if args.flip_vertical is None
691 : : else args.flip_vertical
692 : : ),
693 : : )
694 : :
695 : 1 : output = args.output or family.default_output
696 [ - + ]: 1 : if not output.is_absolute():
697 : 0 : output = Path(
698 : : os.environ.get("BUILD_WORKSPACE_DIRECTORY", ".")
699 : : ).joinpath(output)
700 : 1 : generate(output, family, cfg)
701 : :
702 : :
703 [ - + ]: 1 : if __name__ == "__main__":
704 : 0 : main()
|
