From a6bc14f6b40c0626e03297e5ef2aa781c98827c7 Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Tue, 13 Oct 2026 16:13:03 -0700 Subject: [PATCH 1/5] feat(sim): visual randomization of background, materials and lighting Adds VisualRandomizer, which modifies an MjSpec before compilation: - background: attaches an external scene (e.g. a room) as static, collision-free visual geometry and strips ceilings and its lights - material: overwrites materials or assigns new ones to geoms, with optional mesh replacements that carry texture coordinates - lighting: replaces scene lights with sampled presets Each randomizer has an enabled flag that defaults to off. The randomizers never add joints, so the sim state layout stays compatible with recorded episodes. Also adds a helper that writes box-projected UVs for STL meshes. --- python/rcs/sim/visual_randomization.py | 374 ++++++++++++++++++++++ python/tests/test_visual_randomization.py | 216 +++++++++++++ 2 files changed, 590 insertions(+) create mode 100644 python/rcs/sim/visual_randomization.py create mode 100644 python/tests/test_visual_randomization.py 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", " + + + + + + + + + + + + + +""" + +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])), + 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 From 5586a3c73321cc88ea9886c66cb8f606c78109f5 Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Tue, 13 Oct 2026 16:13:34 -0700 Subject: [PATCH 2/5] feat(envs): optional visual randomization in sim env config SimEnvCreatorConfig gets a visual_randomization field. When set, the enabled randomizers are applied to the composed scene in create_model, seeded by the config seed. The sampled parameters are kept on the ModelComposer for inspection. --- python/rcs/envs/scenes.py | 9 ++++++ python/rcs/sim/composer.py | 4 ++- python/tests/test_visual_randomization.py | 35 +++++++++++++++++++++++ 3 files changed, 47 insertions(+), 1 deletion(-) 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/tests/test_visual_randomization.py b/python/tests/test_visual_randomization.py index b2d3e5e8..612f82a6 100644 --- a/python/tests/test_visual_randomization.py +++ b/python/tests/test_visual_randomization.py @@ -214,3 +214,38 @@ def test_box_projected_uv_obj_is_textured(tmp_path: Path): ) 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) From c4396d25260f9fd22706f1b5c07fe2b13c36a2fc Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Tue, 13 Oct 2026 16:16:03 -0700 Subject: [PATCH 3/5] feat(replay): visually randomized copies of replayed datasets replay() accepts a list of ReplayVariants. Each variant names the randomizers applied together and a multiplier for the number of copies per episode, where a fractional part selects a deterministic subset of episodes. keep_original controls whether the unrandomized episodes are replayed as well, so e.g. background:1 and background+material:1 yield three times the source data. Each copy gets a seed derived from the base seed, episode uuid, variant and copy index. With variants, every recorded step stores the variant, seed and sampled parameters in info["visual_randomization"]. Without variants the replay behaves as before. --- python/rcs/sim/replayer.py | 195 +++++++++++++++++++++++++++++++--- python/tests/test_replayer.py | 165 ++++++++++++++++++++++++++++ 2 files changed, 345 insertions(+), 15 deletions(-) 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/tests/test_replayer.py b/python/tests/test_replayer.py index c3911ae2..764d9b5a 100644 --- a/python/tests/test_replayer.py +++ b/python/tests/test_replayer.py @@ -1,21 +1,37 @@ +import json import xml.etree.ElementTree as ET from pathlib import Path from typing import Any, cast import duckdb +import gymnasium as gym import numpy as np +import pytest from rcs._core.sim import SimConfig from rcs.envs.base import RelativeTo from rcs.envs.configs import EmptyWorldFR3Duo +from rcs.envs.scenes import SimEnvCreatorConfig from rcs.envs.storage_wrapper import StorageWrapper from rcs.envs.tasks import PickTaskConfig from rcs.sim.replayer import ( + ReplayVariant, load_distinct_uuids, load_trajectory, + replay, replay_trajectory, restore_sim_step, + variant_copies, ) from rcs.sim.sim import ROOT_RELATIVE_FREE_STATE_ENCODING +from rcs.sim.visual_randomization import ( + LightingPreset, + LightingRandomizerConfig, + LightSpec, + MaterialOption, + MaterialRandomizerConfig, + MaterialTarget, + VisualRandomizerConfig, +) import rcs @@ -424,3 +440,152 @@ def test_replayer_restores_root_relative_free_joint_state_across_root_frame_chan assert np.allclose(actual_joint_qpos, expected_joint_qpos, atol=1e-5, rtol=0) finally: replay_env.close() + + +class _VisuallyRandomizedDuo(EmptyWorldFR3Duo): + def config(self) -> 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")], + ) From 6476e4b65160034aeaf40ab058cfffada65a9e88 Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Tue, 13 Oct 2026 16:16:03 -0700 Subject: [PATCH 4/5] feat(cli): variant, keep-original and seed options for replay Example: rcs replay data out --variant background:1 --variant background+material:0.5 --- python/rcs/__main__.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) 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, ) From 97220d0fc0d98eb6ec69658e33dd85f7a996d48a Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Tue, 13 Oct 2026 16:30:47 -0700 Subject: [PATCH 5/5] style: silence pose shape typing in visual randomization test --- python/tests/test_visual_randomization.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/test_visual_randomization.py b/python/tests/test_visual_randomization.py index 612f82a6..3d0f78f6 100644 --- a/python/tests/test_visual_randomization.py +++ b/python/tests/test_visual_randomization.py @@ -87,7 +87,7 @@ def config(tmp_path: Path) -> VisualRandomizerConfig: scenes=[ BackgroundScene( xml_path=str(background_xml), - scene_to_world=rcs.common.Pose(translation=np.array([1.0, 0.0, 0.0])), + scene_to_world=rcs.common.Pose(translation=np.array([1.0, 0.0, 0.0])), # type: ignore name="room", ) ]