| | """Task generator — produces (TaskSpec, FullLatentState) pairs for episodes.
|
| |
|
| | Supports three modes:
|
| | 1. Select from the pre-defined scenario library.
|
| | 2. Randomly perturb a scenario for domain-randomisation.
|
| | 3. Compose a fully procedural scenario (tissue × modality × difficulty).
|
| | """
|
| |
|
| | from __future__ import annotations
|
| |
|
| | from typing import List, Optional, Tuple
|
| |
|
| | import numpy as np
|
| |
|
| | from models import TaskSpec, tools_for_modality, assays_for_modality
|
| |
|
| | from server.simulator.latent_state import (
|
| | CellPopulation,
|
| | ExperimentProgress,
|
| | FullLatentState,
|
| | GeneProgram,
|
| | LatentBiologicalState,
|
| | ResourceState,
|
| | TechnicalState,
|
| | )
|
| | from .scenarios import SCENARIO_LIBRARY, Scenario
|
| | from .procedural_generator import generate_procedural_scenarios
|
| |
|
| |
|
| | class TaskGenerator:
|
| | """Generates task + latent-state pairs for environment episodes."""
|
| |
|
| | def __init__(
|
| | self,
|
| | scenarios: Optional[List[Scenario]] = None,
|
| | domain_randomise: bool = True,
|
| | ):
|
| | if scenarios is not None:
|
| | self.scenarios = scenarios
|
| | else:
|
| | self.scenarios = list(SCENARIO_LIBRARY) + generate_procedural_scenarios(n=20, seed=42)
|
| | self.domain_randomise = domain_randomise
|
| |
|
| | def generate(
|
| | self,
|
| | *,
|
| | seed: Optional[int] = None,
|
| | scenario_name: Optional[str] = None,
|
| | ) -> Tuple[TaskSpec, FullLatentState]:
|
| | rng = np.random.default_rng(seed)
|
| |
|
| | if scenario_name:
|
| | scenario = self._find_scenario(scenario_name)
|
| | else:
|
| | idx = int(rng.integers(0, len(self.scenarios)))
|
| | scenario = self.scenarios[idx]
|
| |
|
| | task = scenario.task.model_copy(deep=True)
|
| | biology = scenario.biology.model_copy(deep=True)
|
| | technical = scenario.technical.model_copy(deep=True)
|
| |
|
| | if self.domain_randomise:
|
| | self._randomise(rng, task, biology, technical)
|
| |
|
| |
|
| | compatible_tools = [t.name for t in tools_for_modality(task.modality)]
|
| | compatible_assays = [a.name for a in assays_for_modality(task.modality)]
|
| | if compatible_tools:
|
| | task.available_tools = compatible_tools
|
| | if compatible_assays:
|
| | task.available_assays = compatible_assays
|
| |
|
| | latent = FullLatentState(
|
| | biology=biology,
|
| | technical=technical,
|
| | progress=ExperimentProgress(),
|
| | resources=ResourceState(
|
| | budget_total=task.budget_limit,
|
| | time_limit_days=task.time_limit_days,
|
| | ),
|
| | hidden_failure_conditions=list(scenario.hidden_failure_conditions),
|
| | task_modality=task.modality,
|
| | rng_seed=seed or 0,
|
| | )
|
| | return task, latent
|
| |
|
| | def list_scenarios(self) -> List[str]:
|
| | return [s.name for s in self.scenarios]
|
| |
|
| |
|
| |
|
| | def _find_scenario(self, name: str) -> Scenario:
|
| | for s in self.scenarios:
|
| | if s.name == name:
|
| | return s
|
| | available = ", ".join(self.list_scenarios())
|
| | raise ValueError(f"Unknown scenario '{name}'. Available: {available}")
|
| |
|
| | def _randomise(
|
| | self,
|
| | rng: np.random.Generator,
|
| | task: TaskSpec,
|
| | bio: LatentBiologicalState,
|
| | tech: TechnicalState,
|
| | ) -> None:
|
| | budget_scale = float(rng.uniform(0.7, 1.3))
|
| | task.budget_limit *= budget_scale
|
| | task.time_limit_days *= float(rng.uniform(0.8, 1.2))
|
| |
|
| | tech.dropout_rate = float(np.clip(
|
| | tech.dropout_rate + rng.normal(0, 0.02), 0.01, 0.3
|
| | ))
|
| | tech.doublet_rate = float(np.clip(
|
| | tech.doublet_rate + rng.normal(0, 0.01), 0.01, 0.15
|
| | ))
|
| | tech.sample_quality = float(np.clip(
|
| | tech.sample_quality + rng.normal(0, 0.05), 0.5, 1.0
|
| | ))
|
| | tech.ambient_rna_fraction = float(np.clip(
|
| | tech.ambient_rna_fraction + rng.normal(0, 0.01), 0.01, 0.15
|
| | ))
|
| | for batch_id in list(tech.batch_effects.keys()):
|
| | tech.batch_effects[batch_id] = float(np.clip(
|
| | tech.batch_effects[batch_id] + rng.normal(0, 0.03), 0.0, 0.4
|
| | ))
|
| |
|
| | for pop in bio.cell_populations:
|
| | pop.proportion = float(np.clip(
|
| | pop.proportion * rng.uniform(0.8, 1.2), 0.01, 0.8
|
| | ))
|
| | total = sum(p.proportion for p in bio.cell_populations) or 1.0
|
| | for pop in bio.cell_populations:
|
| | pop.proportion /= total
|
| |
|
| | for comparison, effects in bio.true_de_genes.items():
|
| | for gene in list(effects.keys()):
|
| | effects[gene] *= float(rng.uniform(0.8, 1.2))
|
| |
|
| | bio.n_true_cells = max(
|
| | 1000,
|
| | int(bio.n_true_cells * rng.uniform(0.6, 1.4)),
|
| | )
|
| |
|
