diff --git a/python/rcs/__main__.py b/python/rcs/__main__.py index 0df6232f..b339d7df 100644 --- a/python/rcs/__main__.py +++ b/python/rcs/__main__.py @@ -20,6 +20,7 @@ camera_specs_to_configs, run_conversion, ) +from rcs.sim.replayer import ReplayVariant from rcs.sim.replayer import replay as replay_dataset from rcs.utils import export_camera_episode_videos, export_episode_videos @@ -92,6 +93,19 @@ def replay( help="Python import statement to execute before resolving the environment. Repeat for multiple imports. Example: --import 'from rcs_duobench.tasks import bin_sort'", ), ] = None, + variants: Annotated[ + list[str] | None, + typer.Option( + "--variant", + help="Visually randomized copies to add, as [+...][:] with randomizers " + "background, material and lighting. The multiplier is the number of copies per episode and may be " + "fractional. Repeat for multiple variants. Example: --variant background:1 --variant background+material:0.5", + ), + ] = None, + keep_original: Annotated[ + bool, typer.Option(help="Whether to also replay the original episodes without visual randomization.") + ] = True, + seed: Annotated[int, typer.Option(help="Seed for the visual randomization.")] = 0, ): _exec_import_statements(imports) replay_dataset( @@ -101,6 +115,9 @@ def replay( frequency=frequency, relative_to=relative_to, env_id=env_id, + variants=[ReplayVariant.parse(variant) for variant in variants or []], + keep_original=keep_original, + seed=seed, ) diff --git a/python/rcs/envs/scenes.py b/python/rcs/envs/scenes.py index 16ca88f9..8d2adcc1 100644 --- a/python/rcs/envs/scenes.py +++ b/python/rcs/envs/scenes.py @@ -5,6 +5,7 @@ from os import PathLike import gymnasium as gym +import numpy as np from gymnasium.envs.registration import EnvCreator from rcs._core.sim import SimCameraConfig, SimConfig, SimGripperConfig, SimRobotConfig from rcs.camera.interface import BaseCameraSet @@ -24,6 +25,7 @@ from rcs.envs.sim import GripperWrapperSim, RobotSimWrapper from rcs.sim.composer import ModelComposer from rcs.sim.sim import Sim +from rcs.sim.visual_randomization import VisualRandomizer, VisualRandomizerConfig import rcs from rcs import GRIPPER_PATHS, SCENE_PATHS, TASKS @@ -132,6 +134,8 @@ class SimEnvCreatorConfig(typing.Generic[TaskConfig]): """dict of camera_name to CameraAdderConfig, cameras will be added to the scene according to the config, camera_name must be unique across all cameras in the scene""" gripper_offsets: dict[str, rcs.common.Pose] | None = None """optional offsets for the gripper from the robot's attachment site""" + visual_randomization: VisualRandomizerConfig | None = None + """optional visual randomization, the enabled randomizers are sampled once when the model is created""" _original_cfg: typing.Any = None """this will hold the original config in case this config has been prefixed""" @@ -322,6 +326,11 @@ def create_model(self, cfg: SimEnvCreatorConfig) -> MjModel: ), ) + if cfg.visual_randomization is not None and cfg.visual_randomization.enabled_randomizers(): + composer.visual_randomization_params = VisualRandomizer(cfg.visual_randomization).apply( + composer.spec, np.random.default_rng(cfg.visual_randomization.seed) + ) + return composer def create_env_from_model(self, cfg: SimEnvCreatorConfig, mjmodel: MjModel) -> gym.Env: diff --git a/python/rcs/sim/composer.py b/python/rcs/sim/composer.py index 48a42155..4ed1fcf6 100644 --- a/python/rcs/sim/composer.py +++ b/python/rcs/sim/composer.py @@ -1,5 +1,5 @@ import os -from typing import Optional +from typing import Any, Optional import mujoco from rcs._core.common import Pose @@ -21,6 +21,8 @@ def __init__( self.add_gravcomp = add_gravcomp self._gravcomp_prefixes: set[str] = set() self._root_relative_replay_free_joints: set[str] = set() + self.visual_randomization_params: dict[str, Any] = {} + """sampled parameters of the visual randomization applied to this scene, if any""" def _resolve_asset_paths(self, spec: mujoco.MjSpec, xml_path: str): """Resolves relative paths to absolute ones.""" diff --git a/python/rcs/sim/replayer.py b/python/rcs/sim/replayer.py index 1f2f3a36..eff0b85f 100644 --- a/python/rcs/sim/replayer.py +++ b/python/rcs/sim/replayer.py @@ -1,3 +1,8 @@ +import copy +import hashlib +import json +import math +from collections.abc import Iterable from dataclasses import dataclass from pathlib import Path from typing import Any @@ -7,9 +12,14 @@ import numpy as np from rcs._core.sim import SimConfig from rcs.envs.base import RelativeTo, SimEnv, SimStateSchema -from rcs.envs.scenes import SimEnvCreator +from rcs.envs.scenes import SimEnvCreator, SimEnvCreatorConfig from rcs.envs.storage_wrapper import StorageWrapper +from rcs.sim.composer import ModelComposer from rcs.sim.sim import RAW_STATE_ENCODING +from rcs.sim.visual_randomization import RANDOMIZER_NAMES, VisualRandomizer + +ORIGINAL_VARIANT = "original" +VISUAL_RANDOMIZATION_INFO_KEY = "visual_randomization" def _normalize_sim_state_schema(value: Any) -> SimStateSchema: @@ -123,6 +133,117 @@ def replay_trajectory(env: gym.Env, recorded_steps: list[RecordedSimStep], headl env.get_wrapper_attr("success")() +@dataclass(frozen=True) +class ReplayVariant: + """A randomized copy of the replayed data. + + Attributes: + randomizers: names of the visual randomizers that are applied together, see ``RANDOMIZER_NAMES``. + multiplier: number of copies per source episode, e.g. 1.0 adds one randomized copy of every episode and + 0.5 adds one copy for a deterministic half of the episodes. + """ + + randomizers: frozenset[str] + multiplier: float = 1.0 + + def __post_init__(self): + if not self.randomizers: + msg = "A variant needs at least one randomizer, the original data is controlled by keep_original." + raise ValueError(msg) + unknown = self.randomizers - set(RANDOMIZER_NAMES) + if unknown: + msg = f"Unknown randomizers {sorted(unknown)}, available: {list(RANDOMIZER_NAMES)}" + raise ValueError(msg) + if self.multiplier < 0: + msg = f"Variant multiplier must be non-negative, got {self.multiplier}." + raise ValueError(msg) + + @property + def name(self) -> str: + return "+".join(sorted(self.randomizers)) + + @classmethod + def parse(cls, value: str) -> "ReplayVariant": + """Parses '[+...][:]', e.g. 'background+material:0.5'.""" + names, _, multiplier = value.partition(":") + return cls( + randomizers=frozenset(name.strip() for name in names.split("+") if name.strip()), + multiplier=float(multiplier) if multiplier else 1.0, + ) + + +def _stable_int(*parts: Any) -> int: + digest = hashlib.sha256(":".join(str(part) for part in parts).encode()).digest() + # 63 bits so that the value fits into a signed int64 parquet column + return int.from_bytes(digest[:8], "little") & ((1 << 63) - 1) + + +def variant_seed(seed: int, uuid: str, variant: ReplayVariant, copy_idx: int) -> int: + """Seed of one randomized copy, stable across runs and independent of the dataset order.""" + return _stable_int(seed, uuid, variant.name, copy_idx) + + +def variant_copies(uuids: list[str], variant: ReplayVariant, seed: int) -> dict[str, int]: + """Number of randomized copies per episode, the fractional part selects a deterministic subset of episodes.""" + whole = math.floor(variant.multiplier) + num_extra = round((variant.multiplier - whole) * len(uuids)) + extra = set(sorted(uuids, key=lambda uuid: _stable_int(seed, variant.name, uuid))[:num_extra]) + return {uuid: whole + (uuid in extra) for uuid in uuids} + + +class VisualRandomizationInfoWrapper(gym.Wrapper): + """Adds the variant, seed and sampled visual randomization parameters to the info dict.""" + + def __init__(self, env: gym.Env, variant: str, seed: int, params: dict[str, Any]): + super().__init__(env) + self._info = {"variant": variant, "seed": seed, "params": json.dumps(params, sort_keys=True)} + + def step(self, action): + obs, reward, terminated, truncated, info = self.env.step(action) + info[VISUAL_RANDOMIZATION_INFO_KEY] = dict(self._info) + return obs, reward, terminated, truncated, info + + def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None): + obs, info = self.env.reset(seed=seed, options=options) + info[VISUAL_RANDOMIZATION_INFO_KEY] = dict(self._info) + return obs, info + + +def _storage_env(env: gym.Env, output: Path | str) -> StorageWrapper: + return StorageWrapper( + env=env, + base_dir=str(output), + instruction=None, + batch_size=32, + max_rows_per_group=100, + max_rows_per_file=1000, + always_record=True, + allow_wrapper_instruction=True, + success_from_env=False, + ) + + +def _replay_episode(env: gym.Env, recorded_steps: list[RecordedSimStep], headless: bool): + env.get_wrapper_attr("set_instruction")(recorded_steps[0].instruction) + replay_trajectory(env, recorded_steps, headless) + + +def _validate_variants(variants: list[ReplayVariant], cfg: SimEnvCreatorConfig): + names = [variant.name for variant in variants] + if len(names) != len(set(names)): + msg = f"Variants must be unique, got {names}. Use the multiplier for more copies of the same variant." + raise ValueError(msg) + requested = set().union(*(variant.randomizers for variant in variants)) + configured = cfg.visual_randomization.configured_randomizers() if cfg.visual_randomization is not None else set() + missing = requested - configured + if missing: + msg = ( + f"Variants request the randomizers {sorted(missing)}, but the environment config does not provide " + "options for them in visual_randomization." + ) + raise ValueError(msg) + + def replay( dataset: Path | str, output: Path | str, @@ -130,7 +251,19 @@ def replay( frequency: int = 30, relative_to: str = RelativeTo.CONFIGURED_ORIGIN.name, env_id: str = "rcs/duo", + variants: Iterable[ReplayVariant] = (), + keep_original: bool = True, + seed: int = 0, ): + """Replays all episodes of a dataset and records them again, optionally with visually randomized copies. + + Without variants, every episode is replayed once with the environment as configured. With variants, the + original episodes (if keep_original) are replayed without visual randomization, and each variant adds + randomized copies, e.g. variants background:1 and background+material:1 result in three times the data. + The randomizer options are taken from the visual_randomization field of the environment config. Every + recorded step then contains the variant, seed and sampled parameters under info["visual_randomization"]. + """ + variants = list(variants) env_creator = gym.envs.registry[env_id].entry_point assert isinstance(env_creator, SimEnvCreator), "only rcs envs are supported for replay" @@ -146,26 +279,58 @@ def replay( if sim_cfg_data.root_frame_objects is None: sim_cfg_data.root_frame_objects = {} + if not variants and not keep_original: + msg = "Nothing to replay: keep_original is disabled and no variants are given." + raise ValueError(msg) + if variants: + _validate_variants(variants, sim_cfg_data) + uuids = load_distinct_uuids(dataset) - env_rel = env_creator.create_env(sim_cfg_data) - env_rel = StorageWrapper( - env=env_rel, - base_dir=str(output), - instruction=None, - batch_size=32, - max_rows_per_group=100, - max_rows_per_file=1000, - always_record=True, - allow_wrapper_instruction=True, - success_from_env=False, + if not variants: + env = _storage_env(env_creator.create_env(sim_cfg_data), output) + try: + for uuid in uuids: + recorded_steps = load_trajectory(dataset, uuid) + if recorded_steps: + _replay_episode(env, recorded_steps, headless) + finally: + env.close() + return + + assert sim_cfg_data.visual_randomization is not None + randomizer = VisualRandomizer(sim_cfg_data.visual_randomization) + # randomization is applied explicitly per variant, the original data stays unrandomized + base_cfg = copy.copy(sim_cfg_data) + base_cfg.visual_randomization = None + copies = {variant: variant_copies(uuids, variant, seed) for variant in variants} + + original_env = ( + _storage_env(VisualRandomizationInfoWrapper(env_creator.create_env(base_cfg), ORIGINAL_VARIANT, 0, {}), output) + if keep_original + else None ) try: for uuid in uuids: recorded_steps = load_trajectory(dataset, uuid) if not recorded_steps: continue - env_rel.get_wrapper_attr("set_instruction")(recorded_steps[0].instruction) - replay_trajectory(env_rel, recorded_steps, headless) + if original_env is not None: + _replay_episode(original_env, recorded_steps, headless) + for variant in variants: + for copy_idx in range(copies[variant][uuid]): + copy_seed = variant_seed(seed, uuid, variant, copy_idx) + composer = env_creator.create_model(base_cfg) + if not isinstance(composer, ModelComposer): + msg = "Visual randomization requires an environment that composes its model with ModelComposer." + raise TypeError(msg) + params = randomizer.apply(composer.spec, np.random.default_rng(copy_seed), set(variant.randomizers)) + env = env_creator.create_env_from_model(base_cfg, composer) + env = _storage_env(VisualRandomizationInfoWrapper(env, variant.name, copy_seed, params), output) + try: + _replay_episode(env, recorded_steps, headless) + finally: + env.close() finally: - env_rel.close() + if original_env is not None: + original_env.close() diff --git a/python/rcs/sim/visual_randomization.py b/python/rcs/sim/visual_randomization.py new file mode 100644 index 00000000..6b8436e2 --- /dev/null +++ b/python/rcs/sim/visual_randomization.py @@ -0,0 +1,374 @@ +"""Visual randomization of composed MuJoCo scenes. + +The randomizers operate on an ``mujoco.MjSpec`` before it is compiled. They only change the visual appearance of +the scene (background, materials, lights) and never add joints, so the simulator state layout stays compatible with +recorded episodes. This makes them usable both when creating an environment and when re-rendering recorded data +with the replayer. +""" + +import os +from dataclasses import dataclass, field +from typing import Any + +import mujoco +import numpy as np +from rcs._core.common import Pose + +BACKGROUND = "background" +MATERIAL = "material" +LIGHTING = "lighting" +RANDOMIZER_NAMES = (BACKGROUND, MATERIAL, LIGHTING) + + +# --- background --- + + +@dataclass(kw_only=True) +class BackgroundScene: + xml_path: str + """MJCF of the scene the setup is placed in, e.g. a room.""" + scene_to_world: Pose = field(default_factory=Pose) + """Pose of the background scene in the mujoco world frame, use it to place the setup on free floor.""" + name: str | None = None + """Optional name used in the sampled parameters, defaults to the xml file name.""" + + +@dataclass(kw_only=True) +class BackgroundRandomizerConfig: + enabled: bool = False + scenes: list[BackgroundScene] = field(default_factory=list) + remove_base_geoms: list[str] = field(default_factory=lambda: ["floor"]) + """Geoms of the base scene that are removed when a background is added, e.g. the default floor plane.""" + remove_name_patterns: list[str] = field(default_factory=lambda: ["ceiling"]) + """Elements of the background scene whose name contains one of these (case-insensitive) substrings are removed. + Ceilings would otherwise block overhead lights.""" + remove_lights: bool = True + """Remove the lights of the background scene so that lighting stays controlled by the main scene.""" + prefix: str = "background_" + + +# --- material --- + + +@dataclass(kw_only=True) +class MaterialOption: + texture: str | None = None + """Path to a 2D texture image, None for an untextured material.""" + rgba: tuple[float, float, float, float] = (1.0, 1.0, 1.0, 1.0) + texrepeat: tuple[float, float] = (1.0, 1.0) + texuniform: bool = False + specular: float = 0.1 + shininess: float = 0.1 + reflectance: float = 0.0 + name: str | None = None + """Optional name used in the sampled parameters, defaults to the texture file name.""" + + +@dataclass(kw_only=True) +class MaterialTarget: + """A group of materials and geoms that receive the same sampled material option.""" + + options: list[MaterialOption] + materials: list[str] = field(default_factory=list) + """Names of existing materials whose properties are overwritten.""" + geoms: list[str] = field(default_factory=list) + """Names of geoms that get a new material assigned.""" + mesh_files: dict[str, str] = field(default_factory=dict) + """Mesh name to replacement file, e.g. a copy with texture coordinates of a mesh that has none.""" + name: str | None = None + + +@dataclass(kw_only=True) +class MaterialRandomizerConfig: + enabled: bool = False + targets: list[MaterialTarget] = field(default_factory=list) + + +# --- lighting --- + + +@dataclass(kw_only=True) +class LightSpec: + type: str = "directional" + """One of 'directional', 'spot' or 'point'.""" + pos: tuple[float, float, float] = (0.0, 0.0, 2.0) + dir: tuple[float, float, float] = (0.0, 0.0, -1.0) + target: tuple[float, float, float] | None = None + """If given, the light points from pos to this point and dir is ignored.""" + diffuse: tuple[float, float, float] = (0.7, 0.7, 0.7) + specular: tuple[float, float, float] = (0.3, 0.3, 0.3) + ambient: tuple[float, float, float] = (0.0, 0.0, 0.0) + cutoff: float = 45.0 + exponent: float = 10.0 + attenuation: tuple[float, float, float] = (1.0, 0.0, 0.0) + castshadow: bool = True + + +@dataclass(kw_only=True) +class LightingPreset: + lights: list[LightSpec] = field(default_factory=list) + headlight_ambient: tuple[float, float, float] | None = None + headlight_diffuse: tuple[float, float, float] | None = None + replace_scene_lights: bool = True + """Remove all lights of the scene before adding the preset lights.""" + name: str | None = None + + +@dataclass(kw_only=True) +class LightingRandomizerConfig: + enabled: bool = False + presets: list[LightingPreset] = field(default_factory=list) + + +@dataclass(kw_only=True) +class VisualRandomizerConfig: + background: BackgroundRandomizerConfig = field(default_factory=BackgroundRandomizerConfig) + material: MaterialRandomizerConfig = field(default_factory=MaterialRandomizerConfig) + lighting: LightingRandomizerConfig = field(default_factory=LightingRandomizerConfig) + shadowsize: int | None = None + """Optional shadow map size, larger values reduce shadow artifacts of spot lights.""" + seed: int | None = None + """Seed used when the randomization is applied at environment creation.""" + + def enabled_randomizers(self) -> set[str]: + return {name for name in RANDOMIZER_NAMES if getattr(self, name).enabled} + + def configured_randomizers(self) -> set[str]: + """Randomizers that have options to sample from, independent of their enabled flag.""" + configured = set() + if self.background.scenes: + configured.add(BACKGROUND) + if self.material.targets: + configured.add(MATERIAL) + if self.lighting.presets: + configured.add(LIGHTING) + return configured + + +def _basename(path: str | None) -> str | None: + return None if path is None else os.path.splitext(os.path.basename(path))[0] + + +def _strip_elements(spec: mujoco.MjSpec, patterns: list[str]): + lowered = [p.lower() for p in patterns] + for collection in (spec.geoms, spec.sites, spec.cameras, spec.lights): + for element in list(collection): + if element.name and any(p in element.name.lower() for p in lowered): + spec.delete(element) + + +def _make_static_visual(spec: mujoco.MjSpec): + """Removes everything that would change the simulator state layout or interact with the setup.""" + for collection in (spec.actuators, spec.sensors, spec.tendons, spec.equalities, spec.keys, spec.excludes): + for element in list(collection): + spec.delete(element) + for joint in list(spec.joints): + spec.delete(joint) + for geom in spec.geoms: + geom.contype = 0 + geom.conaffinity = 0 + + +def apply_background(spec: mujoco.MjSpec, cfg: BackgroundRandomizerConfig, rng: np.random.Generator) -> dict[str, Any]: + if not cfg.scenes: + msg = "Background randomization requires at least one background scene." + raise ValueError(msg) + scene = cfg.scenes[rng.integers(len(cfg.scenes))] + if not os.path.exists(scene.xml_path): + msg = f"Background scene not found: {scene.xml_path}" + raise FileNotFoundError(msg) + + for geom in list(spec.geoms): + if geom.name in cfg.remove_base_geoms: + spec.delete(geom) + + background = mujoco.MjSpec.from_file(scene.xml_path) + xml_dir = os.path.dirname(os.path.abspath(scene.xml_path)) + for mesh in background.meshes: + if mesh.file and not os.path.isabs(mesh.file): + mesh.file = os.path.abspath(os.path.join(xml_dir, background.meshdir or "", mesh.file)) + for tex in background.textures: + if tex.file and not os.path.isabs(tex.file): + tex.file = os.path.abspath(os.path.join(xml_dir, background.texturedir or "", tex.file)) + _strip_elements(background, cfg.remove_name_patterns) + if cfg.remove_lights: + for light in list(background.lights): + background.delete(light) + _make_static_visual(background) + + frame = spec.worldbody.add_frame() + frame.pos = list(scene.scene_to_world.translation()) + frame.quat = list(scene.scene_to_world.rotation_q_wxyz()) + spec.attach(background, prefix=cfg.prefix, suffix="", frame=frame) + return {"scene": scene.name or _basename(scene.xml_path)} + + +def _set_material(spec: mujoco.MjSpec, material: mujoco._specs.MjsMaterial, option: MaterialOption): + textures = list(material.textures) + rgb_role = int(mujoco.mjtTextureRole.mjTEXROLE_RGB) + if option.texture is not None: + texture = spec.add_texture() + texture.name = f"{material.name}_randomized_texture" + texture.type = mujoco.mjtTexture.mjTEXTURE_2D + texture.file = os.path.abspath(option.texture) + textures[rgb_role] = texture.name + else: + textures[rgb_role] = "" + material.textures = textures + material.rgba = list(option.rgba) + material.texrepeat = list(option.texrepeat) + material.texuniform = option.texuniform + material.specular = option.specular + material.shininess = option.shininess + material.reflectance = option.reflectance + + +def apply_material(spec: mujoco.MjSpec, cfg: MaterialRandomizerConfig, rng: np.random.Generator) -> dict[str, Any]: + if not cfg.targets: + msg = "Material randomization requires at least one material target." + raise ValueError(msg) + sampled = {} + geoms_by_name = {geom.name: geom for geom in spec.geoms if geom.name} + for target_idx, target in enumerate(cfg.targets): + target_name = target.name or f"target{target_idx}" + option = target.options[rng.integers(len(target.options))] + for mesh_name, mesh_file in target.mesh_files.items(): + mesh = spec.mesh(mesh_name) + if mesh is None: + msg = f"Mesh '{mesh_name}' of target '{target_name}' not found in scene." + raise ValueError(msg) + mesh.file = os.path.abspath(mesh_file) + for material_name in target.materials: + material = spec.material(material_name) + if material is None: + msg = f"Material '{material_name}' of target '{target_name}' not found in scene." + raise ValueError(msg) + _set_material(spec, material, option) + if target.geoms: + material = spec.add_material() + material.name = f"{target_name}_randomized_material" + _set_material(spec, material, option) + for geom_name in target.geoms: + if geom_name not in geoms_by_name: + msg = f"Geom '{geom_name}' of target '{target_name}' not found in scene." + raise ValueError(msg) + geom = geoms_by_name[geom_name] + geom.material = material.name + # a non-default geom rgba would override the material color + geom.rgba = [0.5, 0.5, 0.5, 1.0] + sampled[target_name] = option.name or _basename(option.texture) or "untextured" + return sampled + + +_LIGHT_TYPES = { + "directional": mujoco.mjtLightType.mjLIGHT_DIRECTIONAL, + "spot": mujoco.mjtLightType.mjLIGHT_SPOT, + "point": mujoco.mjtLightType.mjLIGHT_POINT, +} + + +def apply_lighting(spec: mujoco.MjSpec, cfg: LightingRandomizerConfig, rng: np.random.Generator) -> dict[str, Any]: + if not cfg.presets: + msg = "Lighting randomization requires at least one lighting preset." + raise ValueError(msg) + preset_idx = int(rng.integers(len(cfg.presets))) + preset = cfg.presets[preset_idx] + if preset.replace_scene_lights: + for light in list(spec.lights): + spec.delete(light) + for light_spec in preset.lights: + light = spec.worldbody.add_light() + light.type = _LIGHT_TYPES[light_spec.type] + light.pos = list(light_spec.pos) + direction = np.array(light_spec.dir, dtype=np.float64) + if light_spec.target is not None: + direction = np.subtract(light_spec.target, light_spec.pos, dtype=np.float64) + light.dir = list(direction / np.linalg.norm(direction)) + light.diffuse = list(light_spec.diffuse) + light.specular = list(light_spec.specular) + light.ambient = list(light_spec.ambient) + light.cutoff = light_spec.cutoff + light.exponent = light_spec.exponent + light.attenuation = list(light_spec.attenuation) + light.castshadow = light_spec.castshadow + if preset.headlight_ambient is not None: + spec.visual.headlight.ambient = list(preset.headlight_ambient) + if preset.headlight_diffuse is not None: + spec.visual.headlight.diffuse = list(preset.headlight_diffuse) + return {"preset": preset.name or f"preset{preset_idx}"} + + +class VisualRandomizer: + """Applies the configured randomizers to an MjSpec. + + Example: + randomizer = VisualRandomizer(cfg) + params = randomizer.apply(composer.spec, np.random.default_rng(0)) + """ + + def __init__(self, cfg: VisualRandomizerConfig): + self.cfg = cfg + + def apply( + self, spec: mujoco.MjSpec, rng: np.random.Generator, randomizers: set[str] | None = None + ) -> dict[str, Any]: + """Randomizes the spec in place. + + Args: + spec: scene to randomize, must not be compiled into a running simulation yet. + rng: random generator used for sampling. + randomizers: names of the randomizers to apply, defaults to the enabled ones in the config. + + Returns: + The sampled parameters per randomizer, e.g. {"background": {"scene": "FloorPlan1"}}. + """ + if randomizers is None: + randomizers = self.cfg.enabled_randomizers() + unknown = set(randomizers) - set(RANDOMIZER_NAMES) + if unknown: + msg = f"Unknown randomizers {sorted(unknown)}, available: {list(RANDOMIZER_NAMES)}" + raise ValueError(msg) + if randomizers and self.cfg.shadowsize is not None: + spec.visual.quality.shadowsize = self.cfg.shadowsize + + # fixed order so that the result only depends on the rng seed + sampled: dict[str, Any] = {} + if BACKGROUND in randomizers: + sampled[BACKGROUND] = apply_background(spec, self.cfg.background, rng) + if MATERIAL in randomizers: + sampled[MATERIAL] = apply_material(spec, self.cfg.material, rng) + if LIGHTING in randomizers: + sampled[LIGHTING] = apply_lighting(spec, self.cfg.lighting, rng) + return sampled + + +def write_box_projected_uv_obj(stl_path: str, obj_path: str, scale: float = 0.1) -> str: + """Writes a copy of a binary STL mesh as OBJ with box-projected texture coordinates. + + MuJoCo can only apply 2D textures to meshes with texture coordinates. This helper makes STL meshes usable as + material randomization targets, see ``MaterialTarget.mesh_files``. + + Args: + stl_path: binary STL file. + obj_path: output OBJ file. + scale: texture size in mesh units, i.e. the texture repeats every ``scale`` units. + """ + with open(stl_path, "rb") as f: + data = f.read() + n_tris = int(np.frombuffer(data, np.uint32, 1, 80)[0]) + dtype = np.dtype([("normal", " SimEnvCreatorConfig: + cfg = super().config() + cfg.task_cfg = PickTaskConfig(robot_name="right") + cfg.camera_cfgs = {} + cfg.visual_randomization = VisualRandomizerConfig( + material=MaterialRandomizerConfig( + targets=[ + MaterialTarget( + name="floor", + materials=["groundplane"], + options=[MaterialOption(rgba=(0.8, 0.1, 0.1, 1.0), name="red")], + ) + ] + ), + lighting=LightingRandomizerConfig( + presets=[ + LightingPreset(name="spot", lights=[LightSpec(type="spot", pos=(1, 0, 2), target=(0.5, 0, 0))]) + ] + ), + ) + return cfg + + +VISUALLY_RANDOMIZED_ENV_ID = "rcs/test_visually_randomized_duo" +if VISUALLY_RANDOMIZED_ENV_ID not in gym.envs.registry: + gym.register(id=VISUALLY_RANDOMIZED_ENV_ID, entry_point=_VisuallyRandomizedDuo()) + + +def _visual_randomization_infos(dataset_dir: Path) -> dict[str, dict[str, Any]]: + """Visual randomization info per replayed episode uuid, checks that it is constant within an episode.""" + infos: dict[str, dict[str, Any]] = {} + for uuid in load_distinct_uuids(dataset_dir): + steps = load_trajectory(dataset_dir, uuid) + episode_infos = [step.info["visual_randomization"] for step in steps] + assert all(info == episode_infos[0] for info in episode_infos) + infos[uuid] = episode_infos[0] + return infos + + +def test_replay_variant_parsing(): + variant = ReplayVariant.parse("material+background:0.5") + assert variant.randomizers == frozenset({"background", "material"}) + assert variant.multiplier == 0.5 + assert variant.name == "background+material" + assert ReplayVariant.parse("lighting").multiplier == 1.0 + with pytest.raises(ValueError, match="Unknown randomizers"): + ReplayVariant.parse("camera:1") + with pytest.raises(ValueError, match="at least one randomizer"): + ReplayVariant.parse(":1") + + +def test_variant_copies_are_exact_and_deterministic(): + uuids = [f"episode{i}" for i in range(10)] + copies = variant_copies(uuids, ReplayVariant.parse("lighting:1.3"), seed=0) + assert sum(copies.values()) == 13 + assert set(copies.values()) == {1, 2} + assert copies == variant_copies(list(reversed(uuids)), ReplayVariant.parse("lighting:1.3"), seed=0) + assert sum(variant_copies(uuids, ReplayVariant.parse("lighting:0"), seed=0).values()) == 0 + + +def test_replay_with_variants_adds_randomized_copies(tmp_path: Path): + source_dir = tmp_path / "source" + replay_dir = tmp_path / "replayed" + for _ in range(2): + _record_source_dataset(source_dir, limit=3, instruction="pick up cube") + source_steps = {uuid: load_trajectory(source_dir, uuid) for uuid in load_distinct_uuids(source_dir)} + + replay( + source_dir, + replay_dir, + env_id=VISUALLY_RANDOMIZED_ENV_ID, + variants=[ReplayVariant.parse("lighting:1"), ReplayVariant.parse("lighting+material:0.5")], + ) + + infos = _visual_randomization_infos(replay_dir) + variants = sorted(info["variant"] for info in infos.values()) + assert variants == ["lighting", "lighting", "lighting+material", "original", "original"] + for info in infos.values(): + params = json.loads(info["params"]) + if info["variant"] == "original": + assert params == {} + else: + assert params["lighting"] == {"preset": "spot"} + assert ("material" in params) == ("material" in info["variant"]) + + # randomized copies replay the same simulator states as the source episodes + reference = next(iter(source_steps.values())) + for uuid in load_distinct_uuids(replay_dir): + replay_steps = load_trajectory(replay_dir, uuid) + assert any( + all( + np.allclose(r.sim_state, s.sim_state, atol=1e-5, rtol=0) for r, s in zip(replay_steps, src, strict=True) + ) + for src in source_steps.values() + ) + assert len(replay_steps) == len(reference) + + +def test_replay_without_original_and_multiple_copies(tmp_path: Path): + source_dir = tmp_path / "source" + replay_dir = tmp_path / "replayed" + _record_source_dataset(source_dir, limit=2, instruction="pick up cube") + + replay( + source_dir, + replay_dir, + env_id=VISUALLY_RANDOMIZED_ENV_ID, + variants=[ReplayVariant.parse("lighting:2")], + keep_original=False, + ) + + infos = list(_visual_randomization_infos(replay_dir).values()) + assert [info["variant"] for info in infos] == ["lighting", "lighting"] + assert infos[0]["seed"] != infos[1]["seed"] + + +def test_replay_without_variants_has_no_randomization_info(tmp_path: Path): + source_dir = tmp_path / "source" + replay_dir = tmp_path / "replayed" + _record_source_dataset(source_dir, limit=2, instruction="pick up cube") + + replay(source_dir, replay_dir, env_id=VISUALLY_RANDOMIZED_ENV_ID) + + for uuid in load_distinct_uuids(replay_dir): + assert all("visual_randomization" not in step.info for step in load_trajectory(replay_dir, uuid)) + with pytest.raises(ValueError, match="Nothing to replay"): + replay(source_dir, tmp_path / "empty", env_id=VISUALLY_RANDOMIZED_ENV_ID, keep_original=False) + + +def test_replay_variants_require_configured_randomizers(tmp_path: Path): + source_dir = tmp_path / "source" + _record_source_dataset(source_dir, limit=2, instruction="pick up cube") + with pytest.raises(ValueError, match="background"): + replay( + source_dir, + tmp_path / "replayed", + env_id=VISUALLY_RANDOMIZED_ENV_ID, + variants=[ReplayVariant.parse("background")], + ) + with pytest.raises(ValueError, match="unique"): + replay( + source_dir, + tmp_path / "replayed", + env_id=VISUALLY_RANDOMIZED_ENV_ID, + variants=[ReplayVariant.parse("lighting:1"), ReplayVariant.parse("lighting:0.5")], + ) diff --git a/python/tests/test_visual_randomization.py b/python/tests/test_visual_randomization.py new file mode 100644 index 00000000..3d0f78f6 --- /dev/null +++ b/python/tests/test_visual_randomization.py @@ -0,0 +1,251 @@ +import struct +from pathlib import Path + +import mujoco +import numpy as np +import pytest +from PIL import Image +from rcs.sim.visual_randomization import ( + BACKGROUND, + LIGHTING, + MATERIAL, + BackgroundRandomizerConfig, + BackgroundScene, + LightingPreset, + LightingRandomizerConfig, + LightSpec, + MaterialOption, + MaterialRandomizerConfig, + MaterialTarget, + VisualRandomizer, + VisualRandomizerConfig, + write_box_projected_uv_obj, +) + +import rcs + +BASE_XML = """ + + + + + + + + + + + + + + +""" + +BACKGROUND_XML = """ + + + + + + + + + + + +""" + + +def _write_png(path: Path, color: tuple[int, int, int]) -> str: + Image.new("RGB", (8, 8), color).save(path) + return str(path) + + +def _write_binary_stl(path: Path) -> str: + tris = np.array( + [ + [[0, 0, 0], [1, 0, 0], [0, 1, 0]], + [[0, 0, 0], [0, 1, 0], [0, 0, 1]], + [[0, 0, 0], [0, 0, 1], [1, 0, 0]], + [[1, 0, 0], [0, 0, 1], [0, 1, 0]], + ], + dtype=np.float32, + ) + with open(path, "wb") as f: + f.write(b"\0" * 80 + struct.pack(" VisualRandomizerConfig: + background_xml = tmp_path / "room.xml" + background_xml.write_text(BACKGROUND_XML) + return VisualRandomizerConfig( + background=BackgroundRandomizerConfig( + scenes=[ + BackgroundScene( + xml_path=str(background_xml), + scene_to_world=rcs.common.Pose(translation=np.array([1.0, 0.0, 0.0])), # type: ignore + name="room", + ) + ] + ), + material=MaterialRandomizerConfig( + targets=[ + MaterialTarget( + name="table", + materials=["table_mat"], + options=[MaterialOption(texture=_write_png(tmp_path / "marble.png", (230, 230, 230)))], + ), + MaterialTarget( + name="box", + geoms=["box_geom"], + options=[MaterialOption(rgba=(0.1, 0.2, 0.9, 1.0), name="blue")], + ), + ] + ), + lighting=LightingRandomizerConfig( + presets=[ + LightingPreset( + name="warm_spot", + lights=[LightSpec(type="spot", pos=(1, 0, 2), target=(0, 0, 0.8), diffuse=(1.0, 0.7, 0.4))], + headlight_ambient=(0.1, 0.1, 0.1), + ) + ] + ), + ) + + +def _base_spec() -> mujoco.MjSpec: + return mujoco.MjSpec.from_string(BASE_XML) + + +def test_defaults_are_disabled(): + cfg = VisualRandomizerConfig() + assert cfg.enabled_randomizers() == set() + assert cfg.configured_randomizers() == set() + spec = _base_spec() + assert VisualRandomizer(cfg).apply(spec, np.random.default_rng(0)) == {} + model = spec.compile() + assert model.ngeom == _base_spec().compile().ngeom + + +def test_enabled_flags_select_randomizers(config: VisualRandomizerConfig): + config.lighting.enabled = True + sampled = VisualRandomizer(config).apply(_base_spec(), np.random.default_rng(0)) + assert set(sampled) == {LIGHTING} + assert config.configured_randomizers() == {BACKGROUND, MATERIAL, LIGHTING} + + +def test_background_is_static_and_keeps_state_layout(config: VisualRandomizerConfig): + reference = _base_spec().compile() + spec = _base_spec() + sampled = VisualRandomizer(config).apply(spec, np.random.default_rng(0), {BACKGROUND}) + model = spec.compile() + + assert sampled == {BACKGROUND: {"scene": "room"}} + assert (model.nq, model.nv, model.njnt) == (reference.nq, reference.nv, reference.njnt) + geom_names = {model.geom(i).name for i in range(model.ngeom)} + assert "floor" not in geom_names + assert "background_room_floor" in geom_names + assert "background_ceiling_panel" not in geom_names + assert {model.light(i).name for i in range(model.nlight)} == {"main_light"} + background_geoms = [i for i in range(model.ngeom) if model.geom(i).name.startswith("background_")] + assert all(model.geom_contype[i] == 0 and model.geom_conaffinity[i] == 0 for i in background_geoms) + data = mujoco.MjData(model) + mujoco.mj_forward(model, data) + assert np.allclose(data.geom("background_room_floor").xpos, [1.0, 0.0, 0.0]) + + +def test_material_randomization(config: VisualRandomizerConfig): + spec = _base_spec() + sampled = VisualRandomizer(config).apply(spec, np.random.default_rng(0), {MATERIAL}) + model = spec.compile() + + assert sampled == {MATERIAL: {"table": "marble", "box": "blue"}} + table_mat = model.material("table_mat") + assert model.mat_texid[table_mat.id, int(mujoco.mjtTextureRole.mjTEXROLE_RGB)] >= 0 + box_geom = model.geom("box_geom") + assert np.allclose(model.mat_rgba[box_geom.matid[0]], [0.1, 0.2, 0.9, 1.0]) + + +def test_lighting_randomization(config: VisualRandomizerConfig): + spec = _base_spec() + sampled = VisualRandomizer(config).apply(spec, np.random.default_rng(0), {LIGHTING}) + model = spec.compile() + + assert sampled == {LIGHTING: {"preset": "warm_spot"}} + assert model.nlight == 1 + assert model.light_type[0] == mujoco.mjtLightType.mjLIGHT_SPOT + expected_dir = np.array([-1.0, 0.0, -1.2]) / np.linalg.norm([-1.0, 0.0, -1.2]) + assert np.allclose(model.light_dir[0], expected_dir) + assert np.allclose(model.vis.headlight.ambient, [0.1, 0.1, 0.1]) + + +def test_sampling_is_deterministic(config: VisualRandomizerConfig): + config.lighting.presets = [LightingPreset(name=f"p{i}") for i in range(10)] + config.material.targets[1].options = [MaterialOption(name=f"m{i}") for i in range(10)] + randomizer = VisualRandomizer(config) + names = {BACKGROUND, MATERIAL, LIGHTING} + first = randomizer.apply(_base_spec(), np.random.default_rng(42), names) + second = randomizer.apply(_base_spec(), np.random.default_rng(42), names) + assert first == second + + +def test_unknown_and_unconfigured_randomizers_raise(): + randomizer = VisualRandomizer(VisualRandomizerConfig()) + with pytest.raises(ValueError, match="Unknown randomizers"): + randomizer.apply(_base_spec(), np.random.default_rng(0), {"camera"}) + with pytest.raises(ValueError, match="at least one"): + randomizer.apply(_base_spec(), np.random.default_rng(0), {LIGHTING}) + + +def test_box_projected_uv_obj_is_textured(tmp_path: Path): + obj = write_box_projected_uv_obj(_write_binary_stl(tmp_path / "part.stl"), str(tmp_path / "part_uv.obj")) + spec = mujoco.MjSpec.from_string( + f""" + + + + + """ + ) + model = spec.compile() + assert model.mesh_texcoordnum[0] == 12 + + +def _empty_world_cfg(visual_randomization: VisualRandomizerConfig | None): + from rcs.envs.configs import EmptyWorldFR3Duo + + scene = EmptyWorldFR3Duo() + cfg = scene.config() + cfg.headless = True + cfg.camera_cfgs = {} + cfg.visual_randomization = visual_randomization + return scene, cfg + + +def test_env_creation_applies_enabled_randomizers(config: VisualRandomizerConfig): + config.lighting.enabled = True + scene, cfg = _empty_world_cfg(config) + composer = scene.create_model(cfg) + assert composer.visual_randomization_params == {LIGHTING: {"preset": "warm_spot"}} + env = scene.create_env_from_model(cfg, composer) + try: + model = env.get_wrapper_attr("sim").model + assert model.nlight == 1 + assert model.light_type[0] == mujoco.mjtLightType.mjLIGHT_SPOT + finally: + env.close() + + +def test_env_creation_without_enabled_randomizers_is_unchanged(config: VisualRandomizerConfig): + scene, cfg = _empty_world_cfg(None) + reference = scene.create_model(cfg).spec.compile() + scene, cfg = _empty_world_cfg(config) + composer = scene.create_model(cfg) + model = composer.spec.compile() + assert composer.visual_randomization_params == {} + assert (model.ngeom, model.nlight, model.nmat) == (reference.ngeom, reference.nlight, reference.nmat)