Skip to content

tinker_cookbook.rl.RLDataset

class tinker_cookbook.rl.RLDataset(ABC)

A dataset that produces batches of EnvGroupBuilder instances.

This is the primary dataset interface consumed by RL training loops. Implementations must define get_batch and __len__.

class MyDataset(RLDataset):
def __init__(self, builders: list[EnvGroupBuilder], batch_size: int):
self.builders = builders
self.batch_size = batch_size
def get_batch(self, index: int) -> Sequence[EnvGroupBuilder]:
start = index * self.batch_size
return self.builders[start : start + self.batch_size]
def __len__(self) -> int:
return len(self.builders) // self.batch_size

get_batch(index)

Return a batch of EnvGroupBuilder instances at the given index.

Parameters:

  • index (int) – The batch index (0 <= index < len(self)).

Returns: Sequence[EnvGroupBuilder] – The environment group builders for this batch.

Abstract method.