Implementation:Online ml River Bandit Datasets Base
| Knowledge Sources | |
|---|---|
| Domains | Online_Learning, Multi_Armed_Bandits, Datasets |
| Last Updated | 2026-02-08 16:00 GMT |
Overview
Base class for bandit datasets that extends River's standard dataset interface with bandit-specific functionality.
Description
BanditDataset is an abstract base class that inherits from River's Dataset class and adds bandit-specific properties. It defines the task type as "BANDIT" and requires subclasses to implement an arms property that returns the list of available arms. The class tracks standard dataset metadata like number of features, samples, and classes, while adding arm information to the representation.
Usage
Use this base class when creating new bandit datasets for River. Subclass BanditDataset and implement the arms property to return the list of available arms in your dataset. This ensures consistency across bandit datasets in the River ecosystem.
Code Reference
Source Location
- Repository: Online_ml_River
- File: river/bandit/datasets/base.py
Signature
class BanditDataset(datasets.base.Dataset):
def __init__(
self,
n_features,
n_samples=None,
n_classes=None,
n_outputs=None,
sparse=False,
):
...
@abc.abstractproperty
def arms(self) -> list[bandit.base.ArmID]:
"""The list of arms that can be pulled."""
Import
from river import bandit
I/O Contract
| Property | Type | Description |
|---|---|---|
| arms | list[ArmID] | List of available arms (must be implemented by subclass) |
| n_features | int | Number of features in dataset |
| n_samples | int (optional) | Total number of samples |
Usage Examples
from river import bandit
from river.bandit.datasets.base import BanditDataset
# Subclass to create a custom bandit dataset
class MyBanditDataset(BanditDataset):
def __init__(self):
super().__init__(
n_features=10,
n_samples=1000
)
@property
def arms(self):
return ['arm_1', 'arm_2', 'arm_3']
def _iter(self):
# Implement iteration logic
for i in range(self.n_samples):
context = {f'feature_{j}': j for j in range(10)}
arm = self.arms[i % 3]
reward = 1.0
yield context, arm, reward
# Use the dataset
dataset = MyBanditDataset()
print(dataset.arms)