synth-ai 0.2.4.dev7__py3-none-any.whl → 0.2.4.dev9__py3-none-any.whl

synth_ai/environments/examples/bandit/engine.py

@@ -0,0 +1,294 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.reproducibility.core import IReproducibleEngine
+from synth_ai.environments.stateful.engine import StatefulEngine, StatefulEngineSnapshot
+from synth_ai.environments.tasks.core import TaskInstance
+
+
+@dataclass
+class BanditPublicState:
+    arm_count: int
+    step_count: int
+    max_steps: int
+    last_arm: int | None
+    last_reward: float | None
+    cumulative_reward: float
+    reward_history: list[float]
+    arm_pull_counts: list[int]
+    terminated: bool
+    status: str
+
+    @property
+    def steps_remaining(self) -> int:
+        return max(0, self.max_steps - self.step_count)
+
+    @property
+    def average_reward(self) -> float:
+        return self.cumulative_reward / self.step_count if self.step_count else 0.0
+
+
+@dataclass
+class BanditPrivateState:
+    reward_last: float
+    total_reward: float
+    terminated: bool
+    truncated: bool
+    true_expected_rewards: list[float]
+    step_count: int
+
+
+@dataclass
+class BanditEngineSnapshot(StatefulEngineSnapshot):
+    task_instance_dict: dict[str, Any]
+    engine_snapshot: dict[str, Any]
+
+
+class BanditEngine(StatefulEngine, IReproducibleEngine):
+    def __init__(self, task_instance: TaskInstance):
+        self.task_instance = task_instance
+
+        metadata = getattr(task_instance, "metadata", None)
+        self.bandit_type: str = getattr(metadata, "bandit_type", "bernoulli")
+        self.max_steps: int = int(getattr(metadata, "max_steps", 100))
+        self.seed: int | None = getattr(metadata, "seed", None)
+
+        self.arm_probabilities: list[float] | None = None
+        self.arm_means: list[float] | None = None
+        self.arm_stds: list[float] | None = None
+
+        if self.bandit_type.lower() == "bernoulli":
+            probs = list(getattr(metadata, "arm_probabilities", []) or [])
+            if not probs:
+                probs = [0.1, 0.5, 0.9]
+            self.arm_probabilities = [float(p) for p in probs]
+            self.true_expected_rewards = self.arm_probabilities.copy()
+        elif self.bandit_type.lower() == "gaussian":
+            means = list(getattr(metadata, "arm_means", []) or [])
+            if not means:
+                means = [0.0, 0.5, 1.0]
+            stds = getattr(metadata, "arm_stds", None)
+            if stds is None:
+                stds_list = [0.1] * len(means)
+            elif isinstance(stds, int | float):
+                stds_list = [float(stds)] * len(means)
+            else:
+                stds_list = [float(s) for s in stds]
+            if len(stds_list) != len(means):
+                raise ValueError("arm_stds must match arm_means length")
+            self.arm_means = [float(m) for m in means]
+            self.arm_stds = stds_list
+            self.true_expected_rewards = self.arm_means.copy()
+        else:
+            raise ValueError(f"Unsupported bandit_type: {self.bandit_type}")
+
+        self.arm_count = len(self.true_expected_rewards)
+        if self.arm_count == 0:
+            raise ValueError("Bandit must have at least one arm")
+
+        self._rng = np.random.default_rng(self.seed)
+
+        # Runtime state
+        self.step_count: int = 0
+        self.total_reward: float = 0.0
+        self.last_reward: float = 0.0
+        self.last_arm: int | None = None
+        self.reward_history: list[float] = []
+        self.arm_history: list[int] = []
+        self.arm_pull_counts: list[int] = [0 for _ in range(self.arm_count)]
+        self.terminated: bool = False
+        self.status: str = "in_progress"
+
+    async def _reset_engine(self) -> tuple[BanditPrivateState, BanditPublicState]:
+        if self.seed is not None:
+            self._rng = np.random.default_rng(self.seed)
+        else:
+            self._rng = np.random.default_rng()
+
+        self.step_count = 0
+        self.total_reward = 0.0
+        self.last_reward = 0.0
+        self.last_arm = None
+        self.reward_history = []
+        self.arm_history = []
+        self.arm_pull_counts = [0 for _ in range(self.arm_count)]
+        self.terminated = False
+        self.status = "in_progress"
+
+        private_state = self._build_private_state(reward=0.0)
+        public_state = self._build_public_state(reward=None)
+        return private_state, public_state
+
+    async def _step_engine(self, arm_index: int) -> tuple[BanditPrivateState, BanditPublicState]:
+        if self.terminated:
+            raise RuntimeError("Bandit episode already terminated")
+
+        if not isinstance(arm_index, int):
+            raise TypeError("arm_index must be an integer")
+        if arm_index < 0 or arm_index >= self.arm_count:
+            raise ValueError(f"arm_index {arm_index} out of range 0..{self.arm_count - 1}")
+
+        reward = float(self._sample_reward(arm_index))
+
+        self.step_count += 1
+        self.total_reward += reward
+        self.last_reward = reward
+        self.last_arm = arm_index
+        self.reward_history.append(reward)
+        self.arm_history.append(arm_index)
+        self.arm_pull_counts[arm_index] += 1
+
+        if self.step_count >= self.max_steps:
+            self.terminated = True
+            self.status = "completed"
+
+        private_state = self._build_private_state(reward=reward)
+        public_state = self._build_public_state(reward=reward)
+        return private_state, public_state
+
+    def _sample_reward(self, arm_index: int) -> float:
+        if self.bandit_type.lower() == "bernoulli":
+            assert self.arm_probabilities is not None
+            success_prob = self.arm_probabilities[arm_index]
+            return 1.0 if self._rng.random() < success_prob else 0.0
+
+        if self.bandit_type.lower() == "gaussian":
+            assert self.arm_means is not None and self.arm_stds is not None
+            mean = self.arm_means[arm_index]
+            std = self.arm_stds[arm_index]
+            return float(self._rng.normal(loc=mean, scale=std))
+
+        raise RuntimeError(f"Unknown bandit_type during sampling: {self.bandit_type}")
+
+    def get_current_states_for_observation(
+        self,
+    ) -> tuple[BanditPrivateState, BanditPublicState]:
+        private_state = self._build_private_state(reward=self.last_reward)
+        public_state = self._build_public_state(reward=self.last_reward)
+        return private_state, public_state
+
+    def _build_private_state(self, reward: float) -> BanditPrivateState:
+        return BanditPrivateState(
+            reward_last=float(reward),
+            total_reward=float(self.total_reward),
+            terminated=self.terminated,
+            truncated=False,
+            true_expected_rewards=self.true_expected_rewards.copy(),
+            step_count=self.step_count,
+        )
+
+    def _build_public_state(self, reward: float | None) -> BanditPublicState:
+        return BanditPublicState(
+            arm_count=self.arm_count,
+            step_count=self.step_count,
+            max_steps=self.max_steps,
+            last_arm=self.last_arm,
+            last_reward=float(reward) if reward is not None else (self.last_reward if self.step_count else None),
+            cumulative_reward=float(self.total_reward),
+            reward_history=self.reward_history.copy(),
+            arm_pull_counts=self.arm_pull_counts.copy(),
+            terminated=self.terminated,
+            status=self.status,
+        )
+
+    async def _serialize_engine(self) -> BanditEngineSnapshot:
+        snapshot = {
+            "bandit_type": self.bandit_type,
+            "arm_probabilities": self.arm_probabilities,
+            "arm_means": self.arm_means,
+            "arm_stds": self.arm_stds,
+            "max_steps": self.max_steps,
+            "seed": self.seed,
+            "step_count": self.step_count,
+            "total_reward": self.total_reward,
+            "last_reward": self.last_reward,
+            "last_arm": self.last_arm,
+            "reward_history": self.reward_history,
+            "arm_history": self.arm_history,
+            "arm_pull_counts": self.arm_pull_counts,
+            "terminated": self.terminated,
+            "status": self.status,
+            "true_expected_rewards": self.true_expected_rewards,
+            "rng_state": self._rng.bit_generator.state,
+        }
+        return BanditEngineSnapshot(
+            task_instance_dict=await self.task_instance.serialize(),
+            engine_snapshot=snapshot,
+        )
+
+    @classmethod
+    async def _deserialize_engine(cls, snapshot: BanditEngineSnapshot) -> BanditEngine:
+        from .taskset import BanditTaskInstance
+
+        task_instance = await BanditTaskInstance.deserialize(snapshot.task_instance_dict)
+        engine = cls(task_instance)
+
+        data = snapshot.engine_snapshot
+        engine.bandit_type = data.get("bandit_type", engine.bandit_type)
+        engine.max_steps = int(data.get("max_steps", engine.max_steps))
+        engine.seed = data.get("seed", engine.seed)
+        engine.arm_probabilities = data.get("arm_probabilities", engine.arm_probabilities)
+        engine.arm_means = data.get("arm_means", engine.arm_means)
+        engine.arm_stds = data.get("arm_stds", engine.arm_stds)
+        engine.true_expected_rewards = list(data.get("true_expected_rewards", engine.true_expected_rewards))
+        engine.arm_count = len(engine.true_expected_rewards)
+
+        engine.step_count = int(data.get("step_count", 0))
+        engine.total_reward = float(data.get("total_reward", 0.0))
+        engine.last_reward = float(data.get("last_reward", 0.0))
+        engine.last_arm = data.get("last_arm")
+        engine.reward_history = list(data.get("reward_history", []))
+        engine.arm_history = list(data.get("arm_history", []))
+        engine.arm_pull_counts = list(data.get("arm_pull_counts", [0 for _ in range(engine.arm_count)]))
+        engine.terminated = bool(data.get("terminated", False))
+        engine.status = data.get("status", "in_progress")
+
+        engine._rng = np.random.default_rng()
+        rng_state = data.get("rng_state")
+        if rng_state is not None:
+            engine._rng.bit_generator.state = rng_state
+
+        return engine
+
+
+class SynthBanditObservationCallable(GetObservationCallable):
+    async def get_observation(
+        self, pub: BanditPublicState, priv: BanditPrivateState
+    ) -> InternalObservation:
+        observation: InternalObservation = {
+            "arm_count": pub.arm_count,
+            "steps_taken": pub.step_count,
+            "steps_remaining": pub.steps_remaining,
+            "last_arm": pub.last_arm,
+            "last_reward": pub.last_reward,
+            "cumulative_reward": priv.total_reward,
+            "average_reward": pub.average_reward,
+            "arm_pull_counts": pub.arm_pull_counts,
+            "reward_history": pub.reward_history,
+            "terminated": pub.terminated,
+            "status": pub.status,
+        }
+        return observation
+
+
+class SynthBanditCheckpointObservationCallable(GetObservationCallable):
+    async def get_observation(
+        self, pub: BanditPublicState, priv: BanditPrivateState
+    ) -> InternalObservation:
+        observation: InternalObservation = {
+            "arm_count": pub.arm_count,
+            "total_reward": priv.total_reward,
+            "steps_taken": pub.step_count,
+            "best_expected_reward": max(priv.true_expected_rewards) if priv.true_expected_rewards else None,
+            "terminated": pub.terminated,
+            "status": pub.status,
+        }
+        return observation

synth_ai/environments/examples/bandit/environment.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field, ValidationError
+
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.environment.tools import (
+    AbstractTool,
+    EnvToolCall,
+    ToolResult,
+)
+from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
+from synth_ai.environments.stateful.core import StatefulEnvironment
+from synth_ai.environments.tasks.core import TaskInstance
+
+from .engine import (
+    BanditEngine,
+    BanditEngineSnapshot,
+    SynthBanditCheckpointObservationCallable,
+    SynthBanditObservationCallable,
+)
+
+
+class BanditActionInput(BaseModel):
+    arm: int = Field(..., ge=0, description="Index of the arm to pull (0-based)")
+
+
+class BanditInteractTool(AbstractTool):
+    name = "pull_arm"
+    description = "Pull a specific bandit arm to receive a stochastic reward."
+    call_schema = BanditActionInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: BanditEngine):
+        self.engine = engine
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        try:
+            parsed = self.call_schema(**call.args)
+        except ValidationError as exc:
+            return ToolResult(ok=False, error=str(exc))
+
+        try:
+            private_state, public_state = await self.engine._step_engine(parsed.arm)
+        except Exception as exc:  # noqa: BLE001 - propagate as tool error
+            priv, pub = self.engine.get_current_states_for_observation()
+            return ToolResult(
+                ok=False,
+                error=str(exc),
+                payload={"private_state": priv, "public_state": pub},
+            )
+
+        return ToolResult(
+            ok=True,
+            payload={"private_state": private_state, "public_state": public_state},
+        )
+
+
+class BanditEnvironment(StatefulEnvironment, ReproducibleEnvironment[BanditEngine]):
+    def __init__(
+        self,
+        task_instance: TaskInstance,
+        custom_step_obs: GetObservationCallable | None = None,
+        custom_ckpt_obs: GetObservationCallable | None = None,
+    ) -> None:
+        self.name = "Bandit"
+        self.task_instance = task_instance
+        self.custom_step_observation_callable = custom_step_obs or SynthBanditObservationCallable()
+        self.custom_checkpoint_observation_callable = (
+            custom_ckpt_obs or SynthBanditCheckpointObservationCallable()
+        )
+        self.engine = BanditEngine(task_instance)
+        self._interact_tool = BanditInteractTool(self.engine)
+
+    async def initialize(self) -> InternalObservation:
+        priv, pub = await self.engine._reset_engine()
+        return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+
+    async def step(self, tool_calls: list[EnvToolCall] | EnvToolCall | Any) -> InternalObservation:
+        validated_call = self.validate_tool_calls(tool_calls)
+        result = await self._interact_tool(validated_call)
+
+        if result.ok and result.payload:
+            priv = result.payload["private_state"]
+            pub = result.payload["public_state"]
+            return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+
+        priv, pub = self.engine.get_current_states_for_observation()
+        return await self._to_observation(
+            priv,
+            pub,
+            self.custom_step_observation_callable,
+            extra_obs={"error": result.error} if result.error else None,
+        )
+
+    async def checkpoint(self) -> InternalObservation:
+        priv, pub = self.engine.get_current_states_for_observation()
+        return await self._to_observation(priv, pub, self.custom_checkpoint_observation_callable)
+
+    async def terminate(self) -> InternalObservation:
+        priv, pub = self.engine.get_current_states_for_observation()
+        self.engine.terminated = True
+        pub.terminated = True
+        priv.terminated = True
+        if pub.status != "completed":
+            pub.status = "terminated"
+        return await self._to_observation(priv, pub, self.custom_checkpoint_observation_callable)
+
+    def validate_tool_calls(self, tool_calls: Any) -> EnvToolCall:
+        if isinstance(tool_calls, EnvToolCall):
+            candidate = tool_calls
+        elif isinstance(tool_calls, dict):
+            if "tool" in tool_calls:
+                candidate = EnvToolCall(tool=tool_calls["tool"], args=tool_calls.get("args", {}))
+            elif "name" in tool_calls:
+                candidate = EnvToolCall(
+                    tool=tool_calls["name"], args=tool_calls.get("parameters", {})
+                )
+            elif "function" in tool_calls:
+                func_spec = tool_calls["function"]
+                candidate = EnvToolCall(
+                    tool=func_spec.get("name", "pull_arm"),
+                    args=func_spec.get("arguments", {}),
+                )
+            else:
+                candidate = EnvToolCall(tool="pull_arm", args=tool_calls)
+        elif isinstance(tool_calls, list):
+            if not tool_calls:
+                raise ValueError("Empty tool calls list")
+            return self.validate_tool_calls(tool_calls[0])
+        else:
+            if isinstance(tool_calls, int):
+                return EnvToolCall(tool="pull_arm", args={"arm": tool_calls})
+            try:
+                return EnvToolCall(tool="pull_arm", args={"arm": int(tool_calls)})
+            except (TypeError, ValueError) as exc:
+                raise TypeError("Unsupported tool_calls format for bandit environment") from exc
+
+        tool_name = candidate.tool
+        if tool_name not in {"pull_arm", "interact"}:
+            raise ValueError(f"Unknown tool: {tool_name}")
+
+        args = dict(candidate.args or {})
+        if "arm" not in args:
+            for alias in ("action", "index", "arm_index", "armId", "choice"):
+                if alias in args:
+                    args["arm"] = args[alias]
+                    break
+        if "arm" not in args:
+            raise ValueError("Missing required 'arm' argument")
+
+        try:
+            arm_idx = int(args["arm"])
+        except (TypeError, ValueError) as exc:
+            raise ValueError("'arm' argument must be an integer") from exc
+
+        if arm_idx < 0:
+            raise ValueError("'arm' argument must be non-negative")
+
+        return EnvToolCall(tool="pull_arm", args={"arm": arm_idx})
+
+    async def _to_observation(
+        self,
+        priv,
+        pub,
+        obs_cb: GetObservationCallable | None,
+        extra_obs: dict[str, Any] | None = None,
+    ) -> InternalObservation:
+        observation: InternalObservation
+        if obs_cb:
+            observation = await obs_cb.get_observation(pub, priv)
+        else:
+            observation = {}
+        if extra_obs and isinstance(observation, dict):
+            observation.update(extra_obs)
+        return observation
+
+    async def _serialize_engine(self) -> BanditEngineSnapshot:
+        return await self.engine._serialize_engine()
+
+    @classmethod
+    async def _deserialize_engine(
+        cls,
+        snapshot: BanditEngineSnapshot,
+        task_instance: TaskInstance,
+    ) -> BanditEnvironment:
+        env = cls(task_instance)
+        env.engine = await BanditEngine._deserialize_engine(snapshot)
+        env._interact_tool = BanditInteractTool(env.engine)
+        return env

synth_ai/environments/examples/bandit/taskset.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+from uuid import UUID, uuid4
+
+from synth_ai.environments.tasks.core import (
+    Impetus,
+    Intent,
+    SplitInfo,
+    TaskInstance,
+    TaskInstanceMetadata,
+    TaskInstanceSet,
+)
+
+
+@dataclass
+class BanditTaskInstanceMetadata(TaskInstanceMetadata):
+    name: str
+    bandit_type: str
+    arm_probabilities: list[float] | None
+    arm_means: list[float] | None
+    arm_stds: list[float] | None
+    max_steps: int
+    seed: int | None = None
+
+
+@dataclass
+class BanditTaskInstance(TaskInstance):
+    async def serialize(self) -> dict[str, Any]:
+        return {
+            "id": str(self.id),
+            "impetus": {"instructions": self.impetus.instructions},
+            "intent": {
+                "rubric": self.intent.rubric,
+                "gold_trajectories": self.intent.gold_trajectories,
+                "gold_state_diff": self.intent.gold_state_diff,
+            },
+            "metadata": {
+                "name": self.metadata.name,
+                "bandit_type": self.metadata.bandit_type,
+                "arm_probabilities": self.metadata.arm_probabilities,
+                "arm_means": self.metadata.arm_means,
+                "arm_stds": self.metadata.arm_stds,
+                "max_steps": self.metadata.max_steps,
+                "seed": self.metadata.seed,
+            },
+            "is_reproducible": self.is_reproducible,
+            "initial_engine_snapshot": self.initial_engine_snapshot,
+        }
+
+    @classmethod
+    async def deserialize(cls, data: dict[str, Any]) -> BanditTaskInstance:
+        metadata = BanditTaskInstanceMetadata(
+            name=data["metadata"].get("name", "bandit"),
+            bandit_type=data["metadata"].get("bandit_type", "bernoulli"),
+            arm_probabilities=data["metadata"].get("arm_probabilities"),
+            arm_means=data["metadata"].get("arm_means"),
+            arm_stds=data["metadata"].get("arm_stds"),
+            max_steps=int(data["metadata"].get("max_steps", 100)),
+            seed=data["metadata"].get("seed"),
+        )
+        return cls(
+            id=UUID(data["id"]),
+            impetus=Impetus(instructions=data["impetus"]["instructions"]),
+            intent=Intent(
+                rubric=data["intent"].get("rubric", {}),
+                gold_trajectories=data["intent"].get("gold_trajectories"),
+                gold_state_diff=data["intent"].get("gold_state_diff", {}),
+            ),
+            metadata=metadata,
+            is_reproducible=bool(data.get("is_reproducible", True)),
+            initial_engine_snapshot=data.get("initial_engine_snapshot"),
+        )
+
+
+def _expected_rewards(metadata: BanditTaskInstanceMetadata) -> list[float]:
+    if metadata.bandit_type.lower() == "bernoulli" and metadata.arm_probabilities:
+        return [float(p) for p in metadata.arm_probabilities]
+    if metadata.bandit_type.lower() == "gaussian" and metadata.arm_means:
+        return [float(m) for m in metadata.arm_means]
+    return []
+
+
+def _default_bandit_configs() -> list[dict[str, Any]]:
+    return [
+        {
+            "name": "bernoulli-easy",
+            "bandit_type": "bernoulli",
+            "arm_probabilities": [0.1, 0.4, 0.8],
+            "max_steps": 50,
+            "seed": 7,
+        },
+        {
+            "name": "bernoulli-close",
+            "bandit_type": "bernoulli",
+            "arm_probabilities": [0.45, 0.5, 0.55],
+            "max_steps": 75,
+            "seed": 21,
+        },
+        {
+            "name": "gaussian-wide",
+            "bandit_type": "gaussian",
+            "arm_means": [0.0, 0.3, 1.0],
+            "arm_stds": [0.2, 0.2, 0.4],
+            "max_steps": 60,
+            "seed": 14,
+        },
+    ]
+
+
+def _build_impetus(metadata: BanditTaskInstanceMetadata, arm_count: int) -> Impetus:
+    bandit_desc = metadata.bandit_type.capitalize()
+    return Impetus(
+        instructions=(
+            f"You are interacting with a {bandit_desc} multi-armed bandit task. "
+            f"There are {arm_count} arms, indexed from 0 to {arm_count - 1}. "
+            f"You may pull up to {metadata.max_steps} arms. "
+            "Your objective is to maximize the cumulative reward by learning which arm has the highest expected payoff."
+        )
+    )
+
+
+def _build_intent(expected_rewards: Iterable[float]) -> Intent:
+    rewards = list(expected_rewards)
+    best_arm = max(range(len(rewards)), key=rewards.__getitem__) if rewards else 0
+    return Intent(
+        rubric={
+            "goal": "Maximize cumulative reward over the episode",
+            "measurement": "Higher total reward indicates better performance",
+        },
+        gold_trajectories=None,
+        gold_state_diff={
+            "best_arm_index": best_arm,
+            "expected_rewards": rewards,
+        },
+    )
+
+
+async def create_bandit_taskset(
+    configs: list[dict[str, Any]] | None = None,
+) -> TaskInstanceSet:
+    configs = configs or _default_bandit_configs()
+
+    instances: list[BanditTaskInstance] = []
+    for config in configs:
+        metadata = BanditTaskInstanceMetadata(
+            name=config.get("name", "bandit"),
+            bandit_type=config.get("bandit_type", "bernoulli"),
+            arm_probabilities=config.get("arm_probabilities"),
+            arm_means=config.get("arm_means"),
+            arm_stds=config.get("arm_stds"),
+            max_steps=int(config.get("max_steps", 100)),
+            seed=config.get("seed"),
+        )
+
+        expected = _expected_rewards(metadata)
+        arm_count = len(expected) if expected else (
+            len(metadata.arm_probabilities or [])
+            or len(metadata.arm_means or [])
+            or 0
+        )
+        if arm_count == 0:
+            arm_count = 1
+
+        impetus = _build_impetus(metadata, arm_count)
+        intent = _build_intent(expected)
+
+        instance = BanditTaskInstance(
+            id=uuid4(),
+            impetus=impetus,
+            intent=intent,
+            metadata=metadata,
+            is_reproducible=True,
+            initial_engine_snapshot=None,
+        )
+        instances.append(instance)
+
+    # Simple deterministic splits by index
+    val_ids = {instances[i].id for i in range(0, len(instances), 3)}
+    test_ids = {instances[i].id for i in range(1, len(instances), 3)}
+
+    split_info = SplitInfo(
+        val_instance_ids=val_ids,
+        test_instance_ids=test_ids,
+        _is_split_defined=bool(instances),
+    )
+
+    return TaskInstanceSet(
+        name="Bandit Example TaskSet",
+        description="Lightweight Bernoulli and Gaussian bandit tasks inspired by OpenAI Gym bandits.",
+        instances=instances,
+        split_info=split_info,
+    )
+
+
+# Alias for convenience
+
+taskset = create_bandit_taskset