"""Execution bridge for the connected Oasis action world model trainer.""" from __future__ import annotations import importlib.util import json import os import queue import subprocess import sys import threading import time from pathlib import Path from typing import Any from adam.config import ConfigManager from adam.executor import ToolCancelled, ToolContext, ToolExecutionError from adam.oasis_dataset import dataset_directories, validate_oasis_dataset from adam.process_control import set_process_tree_paused, terminate_process_tree FORCE_STOP_TIMEOUT_SECONDS = 30 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} def _external_oasis_root(context: ToolContext) -> Path | None: path = context.root / "config" / "external_tools.json" try: payload = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): return None for entry in payload.get("tools", []): if not isinstance(entry, dict): continue if str(entry.get("id", "")) != "external_oasis_game_trainer": continue root = Path(str(entry.get("backend", {}).get("root", ""))).expanduser() return root if root.is_dir() else None return None def _oasis_root(context: ToolContext) -> Path: folders = ConfigManager(context.root).get("tool_folders", {}) raw = str(folders.get("oasis_trainer", "")) if isinstance(folders, dict) else "" root = Path(raw).expanduser() if raw else (_external_oasis_root(context) or Path()) script = root / "roblox_action_flow_app.py" if not script.is_file(): raise ToolExecutionError( "Oasis roblox_action_flow_app.py was not found. Connect the Oasis-Game-Trainer folder in Settings." ) return root.resolve() def _safe_model_name(value: str) -> str: name = value.strip() if not name or len(name) > 96 or any(char in name for char in "<>:\"/\\|?*\x00"): raise ToolExecutionError("Choose a short Oasis model name without filesystem-reserved characters.") return name def _valid_oasis_resolution(value: str) -> str: """Validate the frame-size invariant required by the external trainer.""" try: width_text, height_text = str(value).lower().split("x") width, height = int(width_text), int(height_text) except (TypeError, ValueError): raise ToolExecutionError( "Oasis resolution must be WIDTHxHEIGHT, for example 256x144 or 512x288." ) from None if width <= 0 or height <= 0 or width * 9 != height * 16 or width % 16 or height % 16: raise ToolExecutionError( "Oasis resolution must be 16:9 and both dimensions must be divisible by 16 " "(use 256x144 or 512x288)." ) return f"{width}x{height}" def _valid_action_model(path: Path) -> bool: try: info = json.loads((path / "action_flow_model_info.json").read_text(encoding="utf-8")) model_type = info.get("model_type") has_unet = (path / "unet" / "config.json").is_file() has_vae = (path / "vae" / "config.json").is_file() and (path / "vae" / "pytorch_model.bin").is_file() return model_type in { "action_conditioned_rectified_flow_video", "action_conditioned_latent_vae_flow_video", "action_conditioned_temporal_latent_flow", "action_conditioned_temporal_pixel_flow", } and has_unet and (model_type in {"action_conditioned_rectified_flow_video", "action_conditioned_temporal_pixel_flow"} or has_vae) except (OSError, ValueError, TypeError, json.JSONDecodeError): return False def _latest_preview(folder: Path) -> Path | None: try: images = [ path for path in folder.rglob("*") if path.is_file() and path.suffix.casefold() in IMAGE_EXTENSIONS and any(token in path.name.casefold() for token in ("preview", "sample", "epoch")) ] return max(images, key=lambda path: path.stat().st_mtime) if images else None except OSError: return None def _parse_event(context: ToolContext, event: dict[str, Any], epochs: int, output: Path, preview_enabled: bool, preview_every: int, preview_steps: int) -> None: kind = str(event.get("type", "")) if kind == "start": context.log( f"Oasis dataset ready: {event.get('transitions', 0):,} transitions, " f"{event.get('training', 0):,} training, {event.get('validation', 0):,} validation on {event.get('device', 'device')}." ) context.progress(1, "Starting Oasis trainer") elif kind == "progress": update = int(event.get("update", 0) or 0) total = int(event.get("total_updates", 0) or 0) epoch = int(event.get("epoch", 0) or 0) loss = event.get("loss") eta = event.get("eta") message = f"Oasis epoch {epoch} of {epochs}" if total: message += f", step {update:,} of {total:,}" if isinstance(loss, (int, float)): message += f", loss {loss:.4f}" context.progress( max(1, min(99, round(update * 100 / total))) if total else max(1, min(99, round(epoch * 100 / max(epochs, 1)))), message, current_step=update, total_steps=total, epoch=epoch, total_epochs=epochs, eta_seconds=eta if isinstance(eta, (int, float)) else None, unit="step" if total else "epoch", ) elif kind == "validation": validation = event.get("validation", {}) loss = validation.get("total_loss") if isinstance(validation, dict) else None context.log(f"Oasis validation epoch {event.get('epoch')}: loss {loss:.4f}" if isinstance(loss, (int, float)) else f"Oasis validation epoch {event.get('epoch')}.") elif kind == "temporal_validation": metrics = event.get("metrics", {}) rollout = metrics.get("rollout_rgb_l1") if isinstance(metrics, dict) else None context.log( f"Temporal Oasis validation epoch {event.get('epoch')}: " + (f"generated rollout RGB L1 {rollout}." if rollout else "completed.") ) elif kind == "preview" and preview_enabled: path = Path(str(event.get("path", ""))) if event.get("path") else _latest_preview(output) if path: epoch = int(event.get("epoch", 0) or 0) context.preview(path, epoch=epoch, next_epoch=epoch + preview_every, steps=preview_steps, kind="training") elif kind == "saved": context.log(f"Oasis checkpoint saved at epoch {event.get('epoch')}.") elif kind == "best_checkpoint": context.log( f"Oasis selected epoch {event.get('epoch')} as the current best held-out checkpoint." ) elif kind == "recovery_saved": context.log(f"Oasis recovery checkpoint saved at batch {event.get('batch')}.") elif kind == "warning": context.log(str(event.get("message", "Oasis trainer warning."))) elif kind == "complete": context.progress(100, "Oasis training completed") def train_oasis( context: ToolContext, dataset_dir: str | list[str], model_name: str, epochs: int, output_dir: str, resume_from: str = "", resolution: str = "256x144", batch_size: int = 2, learning_rate: float = 0.00002, workers: int = 2, mixed_precision: str = "fp16", gradient_accumulation: int = 1, frame_gap: int = 1, sequence_context: int = 1, model_engine: str = "pixel_flow", context_frames: int = 4, rollout_frames: int = 3, vae_epochs: int = 5, action_aggregation: str = "window", validation_split: float = 0.1, validation_batches: int = 8, save_every: int = 5, best_checkpoint_min_improvement: float = 0.1, preview_enabled: bool = True, preview_every: int = 5, preview_steps: int = 1, seed: int = 1234, base_model: str = "", condition_noise: float = 0.03, temporal_loss_weight: float = 0.1, motion_loss_weight: float = 2.0, action_input_scale: float = 8.0, neutral_action_dropout: float = 0.15, action_contrast_weight: float = 0.35, action_contrast_margin: float = 0.02, contrast_every: int = 4, contrast_samples: int = 2, chunk_size: int = 0, chunk_mode: str = "balanced", chunk_offset: int = 0, replay_older_percent: float = 50.0, include_older_data: bool = True, recovery_minutes: int = 30, benchmark_batches: int = 8, tf32: bool = True, gradient_checkpointing: bool = False, balance_actions: bool = True, ) -> dict[str, object]: root = _oasis_root(context) script = root / "roblox_action_flow_app.py" model_name = _safe_model_name(model_name) resolution = _valid_oasis_resolution(resolution) dataset_value = ( ";".join(str(item) for item in dataset_dir) if isinstance(dataset_dir, list) else str(dataset_dir) ) if not include_older_data: directories = dataset_directories(dataset_value) if directories: dataset_value = str(directories[0]) output = Path(output_dir).expanduser().resolve() output_root = (root / "output_action_flow_models").resolve() try: output.relative_to(output_root) except ValueError as exc: raise ToolExecutionError("Oasis outputs must stay inside output_action_flow_models.") from exc resume = Path(resume_from).expanduser().resolve() if resume_from else None if resume: try: output.relative_to(resume) overlaps_resume = True except ValueError: try: resume.relative_to(output) overlaps_resume = True except ValueError: overlaps_resume = False if overlaps_resume: raise ToolExecutionError("Oasis continuation output must be a new folder, separate from the source model.") if output.exists(): raise ToolExecutionError("The chosen Oasis output folder already exists; ADAM will not overwrite it.") if model_engine not in {"pixel_flow", "vae_cpu_lite", "temporal_latent", "temporal_pixel_flow"}: raise ToolExecutionError("Choose pixel_flow, vae_cpu_lite, temporal_latent, or temporal_pixel_flow for Oasis.") if mixed_precision not in {"fp16", "no", "fp32"}: raise ToolExecutionError("Oasis precision must be fp16 or no.") if mixed_precision == "fp32": mixed_precision = "no" if chunk_mode not in {"balanced", "random", "sequential"}: raise ToolExecutionError("Oasis chunk selection must be balanced, random, or sequential.") if model_engine in {"temporal_latent", "temporal_pixel_flow"}: if int(context_frames) not in {1, 4, 8} or int(rollout_frames) not in {1, 3}: raise ToolExecutionError("Temporal Oasis uses 1, 4, or 8 history frames and 1 or 3 future frames.") # The external trainer owns the timing contract; use its actual sequence # inspection instead of accepting old endpoint-only validation here. if str(root) not in sys.path: sys.path.insert(0, str(root)) try: from temporal_oasis import read_sequences from roblox_action_flow_app import BINARY_ACTION_NAMES train_windows, val_windows, temporal_report = read_sequences( dataset_directories(dataset_value), BINARY_ACTION_NAMES, int(context_frames), int(rollout_frames), float(validation_split), ) except (OSError, ValueError, ImportError) as exc: raise ToolExecutionError(f"Temporal Oasis dataset validation failed: {exc}") from exc context.log( f"Temporal Oasis dataset validated: {len(train_windows):,} training windows, " f"{len(val_windows):,} validation windows, {temporal_report['context_frames']} frames of history." ) else: report = validate_oasis_dataset(dataset_value, frame_gap=int(frame_gap)) if not report.ok: raise ToolExecutionError("Oasis dataset validation failed: " + " ".join(report.errors[:12])) context.log( f"Oasis dataset validated: {report.valid_rows:,} labelled frames, " f"{report.valid_transitions:,} transitions, resolution {report.resolution}." ) missing_packages = [ package for package in ("torch", "torchvision", "diffusers", "PIL") if importlib.util.find_spec(package) is None ] if missing_packages: raise ToolExecutionError( "ADAM's Python environment is missing Oasis packages: " + ", ".join(missing_packages) + ". Install the Oasis requirements, then restart ADAM." ) output.parent.mkdir(parents=True, exist_ok=True) if resume and not _valid_action_model(resume): raise ToolExecutionError("Choose a valid Oasis action model folder to resume.") base = Path(base_model).expanduser().resolve() if base_model and model_engine == "pixel_flow" else None command = [ sys.executable, str(script), "--train-worker", "--dataset-dir", dataset_value, "--output-dir", str(output), "--model-name", model_name, "--model-engine", model_engine, "--context-frames", str(int(context_frames)), "--rollout-frames", str(int(rollout_frames)), "--vae-epochs", str(int(vae_epochs)), "--epochs", str(int(epochs)), "--resolution", str(resolution), "--batch-size", str(int(batch_size)), "--learning-rate", str(float(learning_rate)), "--workers", str(int(workers)), "--gradient-accumulation", str(int(gradient_accumulation)), "--condition-noise", str(float(condition_noise)), "--temporal-loss-weight", str(float(temporal_loss_weight)), "--motion-loss-weight", str(float(motion_loss_weight)), "--validation-split", str(float(validation_split)), "--validation-batches", str(int(validation_batches)), "--metrics-every", "10", "--frame-gap", str(int(frame_gap)), "--action-aggregation", str(action_aggregation), "--action-input-scale", str(float(action_input_scale)), "--neutral-action-dropout", str(float(neutral_action_dropout)), "--action-contrast-weight", str(float(action_contrast_weight)), "--action-contrast-margin", str(float(action_contrast_margin)), "--contrast-every", str(max(1, int(contrast_every))), "--contrast-samples", str(max(2, int(contrast_samples))), "--chunk-size", str(max(0, int(chunk_size))), "--chunk-mode", chunk_mode, "--chunk-offset", str(max(0, int(chunk_offset))), "--replay-older-percent", str(max(0.0, float(replay_older_percent))), "--chunk-seed", str(int(seed)), "--benchmark-batches", str(max(1, int(benchmark_batches))), "--recovery-minutes", str(max(0, int(recovery_minutes))), "--save-every", str(int(save_every)), "--best-checkpoint-min-improvement", str(max(0.0, float(best_checkpoint_min_improvement)) / 100.0), "--preview-every", str(int(preview_every) if preview_enabled else int(epochs) + 1), "--preview-steps", str(int(preview_steps)), "--mixed-precision", str(mixed_precision), ] if resume: command.extend(["--continue-action-model", str(resume)]) if base: command.extend(["--base-model", str(base)]) if gradient_checkpointing: command.append("--gradient-checkpointing") if tf32: command.append("--tf32") if balance_actions: command.append("--balance-actions") process = subprocess.Popen( command, cwd=str(root), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding="utf-8", errors="replace", shell=False, creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0, ) lines: queue.Queue[str | None] = queue.Queue() def read_output() -> None: assert process.stdout is not None for line in process.stdout: lines.put(line.rstrip()) lines.put(None) threading.Thread(target=read_output, daemon=True).start() context.progress(1, "Starting Oasis trainer") stopped = False stop_requested_at: float | None = None force_stop_sent = False suspended = False stop_file = output / "stop_action_training.flag" while True: should_pause = not context.run_event.is_set() if should_pause != suspended and set_process_tree_paused(process, should_pause): suspended = should_pause context.log("Oasis trainer paused safely." if suspended else "Oasis trainer resumed.") if context.cancel_event.is_set() and not stopped: if suspended: set_process_tree_paused(process, False) suspended = False output.mkdir(parents=True, exist_ok=True) stop_file.touch(exist_ok=True) stopped = True stop_requested_at = time.monotonic() context.log("Safe stop requested; waiting for Oasis to save a usable checkpoint.") if stop_requested_at and not force_stop_sent and time.monotonic() - stop_requested_at > FORCE_STOP_TIMEOUT_SECONDS: terminate_process_tree(process, timeout=3) force_stop_sent = True context.log("Oasis did not stop in time; terminating the trainer process.") try: line = lines.get(timeout=0.15) if line and line.startswith("ACTION_FLOW_EVENT:"): try: _parse_event(context, json.loads(line.split(":", 1)[1]), int(epochs), output, bool(preview_enabled), int(preview_every), int(preview_steps)) except json.JSONDecodeError: context.log(line) elif line: context.log(line) except queue.Empty: pass if process.poll() is not None and lines.empty(): break if stopped: raise ToolCancelled("Oasis training stopped by user.") if process.returncode != 0: raise ToolExecutionError(f"Oasis trainer exited with code {process.returncode}. See the job log for details.") if not _valid_action_model(output): raise ToolExecutionError("Oasis training ended, but no valid action model metadata was written.") context.progress(100, "Oasis training completed") return { "output_folder": str(output), "model_name": model_name, "assets": [{ "kind": "model", "name": model_name, "path": str(output), "trainer": "oasis", "dataset_path": dataset_value, "checkpoint": str(output), "epochs": int(epochs), }], } def launch_oasis_player( context: ToolContext, model_path: str, model_name: str = "", starting_frame: str = "", seed: int = 0, ) -> dict[str, object]: root = _oasis_root(context) model = Path(model_path).expanduser().resolve() if not _valid_action_model(model): raise ToolExecutionError("Choose a valid Oasis action model folder.") best_checkpoint = model / "best_validation_checkpoint" if _valid_action_model(best_checkpoint): context.log("Using the model's held-out best-validation checkpoint for playback.") model = best_checkpoint if starting_frame: start = Path(starting_frame).expanduser() if not start.is_file(): raise ToolExecutionError("The selected Oasis starting frame does not exist.") command = [sys.executable, str(root / "roblox_action_flow_app.py")] env = os.environ.copy() env["ADAM_OASIS_MODEL_PATH"] = str(model) env["ADAM_OASIS_STARTING_FRAME"] = starting_frame env["ADAM_OASIS_SEED"] = str(int(seed)) process = subprocess.Popen( command, cwd=str(root), env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, shell=False, creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0, ) context.log(f"Launched Oasis playable inference for {model_name or model.name}.") context.progress(100, "Oasis player launched") return {"player_pid": process.pid, "model_path": str(model)}