src/task_set.py

# task_set.py

# 200 pretraining tasks (the 10 held-out unseen tasks are in UNSEEN_TASK_SET below)
TASK_SET = [
    # dmcontrol (21 tasks)
    'walker-stand', 'walker-walk', 'walker-run', 'cheetah-run', 'reacher-easy',
    'reacher-hard', 'acrobot-swingup', 'pendulum-swingup', 'cartpole-balance', 'cartpole-balance-sparse',
    'cartpole-swingup', 'cartpole-swingup-sparse', 'cup-catch', 'finger-spin', 'finger-turn-easy',
    'finger-turn-hard', 'fish-swim', 'hopper-stand', 'hopper-hop', 'quadruped-walk',
    'quadruped-run',
    # dmcontrol-ext (16 tasks)
    'walker-walk-backward', 'walker-run-backward', 'cheetah-run-backward', 'cheetah-run-front', 'cheetah-run-back',
    'cheetah-jump', 'hopper-hop-backward', 'reacher-three-easy', 'reacher-three-hard', 'cup-spin',
    'pendulum-spin', 'jumper-jump', 'spinner-spin', 'spinner-spin-backward', 'spinner-jump',
    'giraffe-run',
    # meta-world (49 tasks)
    'mw-assembly', 'mw-basketball', 'mw-button-press-topdown', 'mw-button-press-topdown-wall', 'mw-button-press',
    'mw-button-press-wall', 'mw-coffee-button', 'mw-coffee-pull', 'mw-coffee-push', 'mw-dial-turn',
    'mw-disassemble', 'mw-door-open', 'mw-door-close', 'mw-drawer-close', 'mw-drawer-open',
    'mw-faucet-open', 'mw-faucet-close', 'mw-hammer', 'mw-handle-press-side', 'mw-handle-press',
    'mw-handle-pull-side', 'mw-handle-pull', 'mw-lever-pull', 'mw-peg-insert-side', 'mw-peg-unplug-side',
    'mw-pick-out-of-hole', 'mw-pick-place', 'mw-pick-place-wall', 'mw-plate-slide', 'mw-plate-slide-side',
    'mw-plate-slide-back', 'mw-plate-slide-back-side', 'mw-push-back', 'mw-push', 'mw-push-wall',
    'mw-reach', 'mw-reach-wall', 'mw-soccer', 'mw-stick-push', 'mw-stick-pull', 
    'mw-sweep-into', 'mw-sweep', 'mw-window-open', 'mw-window-close', 'mw-bin-picking',
    'mw-box-close', 'mw-door-lock', 'mw-door-unlock', 'mw-hand-insert',
    # maniskill (36 tasks)
    'ms-ant-walk', 'ms-ant-run', 'ms-cartpole-balance', 'ms-cartpole-swingup', 'ms-hopper-stand',
    'ms-hopper-hop', 'ms-pick-cube', 'ms-pick-cube-eepose', 'ms-pick-cube-so', 'ms-poke-cube',
    'ms-push-cube', 'ms-pull-cube', 'ms-pull-cube-tool', 'ms-stack-cube', 'ms-place-sphere',
    'ms-lift-peg', 'ms-pick-apple', 'ms-pick-banana', 'ms-pick-can', 'ms-pick-hammer',
    'ms-pick-fork', 'ms-pick-knife', 'ms-pick-mug', 'ms-pick-orange', 'ms-pick-screwdriver',
    'ms-pick-spoon', 'ms-pick-tennis-ball', 'ms-pick-baseball', 'ms-pick-cube-xarm6', 'ms-pick-sponge',
    'ms-anymal-reach', 'ms-reach', 'ms-reach-eepose', 'ms-reach-xarm6', 'ms-cartpole-balance-sparse',
    'ms-cartpole-swingup-sparse',
    # mujoco (6 tasks)
    'mujoco-ant', 'mujoco-halfcheetah', 'mujoco-hopper', 'mujoco-inverted-pendulum', 'mujoco-reacher',
    'mujoco-walker',
    # box2d (8 tasks)
    'bipedal-walker-flat', 'bipedal-walker-uneven', 'bipedal-walker-rugged', 'bipedal-walker-hills', 'bipedal-walker-obstacles',
    'lunarlander-land', 'lunarlander-hover', 'lunarlander-takeoff',
    # robodesk (6 tasks)
    'rd-push-red', 'rd-push-green', 'rd-push-blue', 'rd-open-slide', 'rd-open-drawer',
    'rd-flat-block-in-bin',
    # ogbench (12 tasks)
    'og-ant', 'og-antball', 'og-point-arena', 'og-point-maze', 'og-point-bottleneck',
    'og-point-circle', 'og-point-spiral', 'og-ant-arena', 'og-ant-maze', 'og-ant-bottleneck',
    'og-ant-circle', 'og-ant-spiral',
    # pygame (19 tasks)
    'pygame-cowboy', 'pygame-coinrun', 'pygame-spaceship', 'pygame-pong', 'pygame-bird-attack',
    'pygame-highway', 'pygame-landing', 'pygame-air-hockey', 'pygame-rocket-collect', 'pygame-chase-evade',
    'pygame-coconut-dodge', 'pygame-cartpole-balance', 'pygame-cartpole-swingup', 'pygame-cartpole-balance-sparse', 'pygame-cartpole-swingup-sparse',
    'pygame-cartpole-tremor', 'pygame-point-maze-var1', 'pygame-point-maze-var2', 'pygame-point-maze-var3',
    # atari (27 tasks)
    'atari-alien', 'atari-assault', 'atari-asterix', 'atari-atlantis', 'atari-bank-heist',
    'atari-battle-zone', 'atari-beamrider', 'atari-boxing', 'atari-chopper-command', 'atari-crazy-climber',
    'atari-double-dunk', 'atari-gopher', 'atari-ice-hockey', 'atari-jamesbond', 'atari-kangaroo',
    'atari-krull', 'atari-ms-pacman', 'atari-name-this-game', 'atari-phoenix', 'atari-pong',
    'atari-road-runner', 'atari-robotank', 'atari-seaquest', 'atari-space-invaders', 'atari-tutankham',
    'atari-upndown', 'atari-yars-revenge',
]


# Domain spans in TASK_SET ordering.  Task-set authoritatively defines the
# 10 domains below; keep in sync with the task blocks above.
DOMAIN_SPANS: list[tuple[str, int, int]] = [
    ('dmcontrol',      0,  21),
    ('dmcontrol-ext', 21,  37),
    ('metaworld',     37,  86),
    ('maniskill',     86, 122),
    ('mujoco',       122, 128),
    ('box2d',        128, 136),
    ('robodesk',     136, 142),
    ('ogbench',      142, 154),
    ('pygame',       154, 173),
    ('atari',        173, 200),
]
DOMAINS: list[str] = [name for name, _, _ in DOMAIN_SPANS]

_TASK_TO_DOMAIN: dict[str, str] = {}
for _name, _lo, _hi in DOMAIN_SPANS:
    for _t in TASK_SET[_lo:_hi]:
        _TASK_TO_DOMAIN[_t] = _name


def task_to_domain(task: str) -> str:
    """Map a task name to its canonical domain string.

    Covers TASK_SET ∪ UNSEEN_TASK_SET. Raises KeyError otherwise — the
    caller should fall back to a default domain or update the table if
    this fires.
    """
    return _TASK_TO_DOMAIN[task]


# ---------------------------------------------------------------------------
# Seen / unseen task sets for the targeted-data-collection evaluation.
#   - SEEN_TASK_SET: subset of TASK_SET (tasks the world model trained on).
#   - UNSEEN_TASK_SET: tasks defined in envs/*.py but NOT in TASK_SET (also
#     enumerated as keys in interactive.TEST_TASK_SET).
# Each matched UNSEEN entry is the in-domain analog of a SEEN entry (same
# domain / control morphology), since out-of-domain tasks transfer poorly
# zero-shot.
#
# IMPORTANT: three matched pairs share IDENTICAL CLIP text embeddings
# (cosine = 1.0000) by design: cup-catch ↔ cup-catch-var1, finger-turn-easy ↔
# finger-turn-easy-var1, pygame-point-maze-var1 ↔ pygame-point-maze-var4. The
# model thus has no language signal to disambiguate them, so any UNSEEN-partner
# gain is attributable to visual adaptation, not language conditioning (and
# such pairs are not "true zero-shot" in the language sense). The three pygame
# tasks dungeon-explorer1/foraging/whirlpool have no SEEN partner.
# ---------------------------------------------------------------------------

SEEN_TASK_SET: list[str] = [
    'cup-catch',                # dmcontrol
    'finger-turn-easy',         # dmcontrol
    'mw-push',                  # metaworld
    'ms-push-cube',             # maniskill
    'lunarlander-hover',        # box2d
    'og-point-maze',            # ogbench
    'og-point-bottleneck',      # ogbench
    'pygame-point-maze-var1',   # pygame
    'pygame-pong',              # pygame
    'pygame-bird-attack',       # pygame
]

UNSEEN_TASK_SET: list[str] = [
    # Matched-pair UNSEEN: each entry has a SEEN partner with shared visual style.
    'cup-catch-var1',           # dmcontrol; SEEN partner: cup-catch        (lang cosine 1.0000)
    'finger-turn-easy-var1',    # dmcontrol; SEEN partner: finger-turn-easy (lang cosine 1.0000)
    'ms-push-banana',           # maniskill; SEEN partner: ms-push-cube     (lang cosine 0.986)
    'og-point-var1',            # ogbench;   SEEN partner: og-point-bottleneck (lang cosine 0.992)
    'og-point-var2',            # ogbench;   SEEN partner: og-point-maze       (lang cosine 0.994)
    'pygame-point-maze-var4',   # pygame;    SEEN partner: pygame-point-maze-var1 (lang cosine 1.0000)
    'pygame-reacher-easy',      # pygame;    SEEN partner: reacher-easy (dmcontrol; same instruction)
    # Completely unseen: no SEEN partner.
    'pygame-dungeon-explorer1',
    'pygame-foraging',
    'pygame-whirlpool',
]

# UNSEEN tasks are not part of the contiguous DOMAIN_SPANS over TASK_SET, so
# extend the lookup table explicitly (matched-per-domain home per the
# inline comments above).
_UNSEEN_TASK_TO_DOMAIN: dict[str, str] = {
    'cup-catch-var1':           'dmcontrol',
    'finger-turn-easy-var1':    'dmcontrol',
    'ms-push-banana':           'maniskill',
    'og-point-var1':            'ogbench',
    'og-point-var2':            'ogbench',
    'pygame-point-maze-var4':   'pygame',
    'pygame-reacher-easy':      'pygame',
    'pygame-dungeon-explorer1': 'pygame',
    'pygame-foraging':          'pygame',
    'pygame-whirlpool':         'pygame',
}
_TASK_TO_DOMAIN.update(_UNSEEN_TASK_TO_DOMAIN)


def compute_task_weights(
    tasks: list[str],
    mode: str,
    *,
    targeted_alpha: float = 0.5,
    targeted_tasks: list[str] | None = None,
) -> list[float]:
    """Build per-task sampling weights for WMDataset.

    Modes (all weights are relative; WMDataset normalizes internally):
      - "valid_starts": returns None-equivalent (signals caller to pass
        task_weights=None for legacy P(task) ∝ total valid_starts).
      - "uniform": every task gets weight 1.0 (equal-per-task sampling).
      - "targeted": split total weight `α` uniformly over the focused task
        set and `1-α` uniformly over the rest, biasing sampling toward the
        targeted tasks while retaining the rest. Defaults: `targeted_alpha=0.5`,
        `targeted_tasks=SEEN_TASK_SET ∪ UNSEEN_TASK_SET`.

    The `tasks` list is what the dataset actually loaded (post filter);
    unknown tasks get weight 1.0 under "uniform" (shouldn't happen with
    the canonical TASK_SET).
    """
    if mode == "valid_starts":
        # Sentinel: caller should pass task_weights=None.
        return [1.0] * len(tasks)
    if mode == "uniform":
        return [1.0] * len(tasks)
    if mode == "targeted":
        if not (0.0 <= targeted_alpha <= 1.0):
            raise ValueError(
                f"targeted_alpha must be in [0,1], got {targeted_alpha!r}"
            )
        targeted_set = set(SEEN_TASK_SET + UNSEEN_TASK_SET if targeted_tasks is None else targeted_tasks)
        n_targeted = sum(1 for t in tasks if t in targeted_set)
        n_non = len(tasks) - n_targeted
        # Auto-collapse alpha when one side is empty so we degrade
        # gracefully to "uniform over what's loaded" instead of all-zero
        # weights (which would crash WMDataset's sampler).
        if n_targeted == 0:
            alpha_eff = 0.0
        elif n_non == 0:
            alpha_eff = 1.0
        else:
            alpha_eff = float(targeted_alpha)
        w_targeted = (alpha_eff / n_targeted) if n_targeted > 0 else 0.0
        w_non = ((1.0 - alpha_eff) / n_non) if n_non > 0 else 0.0
        return [w_targeted if t in targeted_set else w_non for t in tasks]
    raise ValueError(f"Unknown task_weighting mode: {mode!r}")