feat: visual randomization for sim scenes and replay

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 <randomizer>[+<randomizer>...][:<multiplier>] 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,
     )
 
 

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:

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."""

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,14 +133,137 @@ 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 '<randomizer>[+<randomizer>...][:<multiplier>]', 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,
     headless: bool = True,
     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()