Add probability dataset (initial: Coin Flip dataset + curriculum)

[kumaranant1]

Hi all 👋,
I’m Anant, and this is my very first open-source contribution (so please excuse any rough edges 🙏🥺).
I’ve been interested in reasoning tasks for a while, and I thought a probability tasks would be a useful addition to reasoning_gym.

As a first step, I’ve implemented the Coin Flip dataset + curriculum as a simple example.
I’d be very happy to receive feedback and improve this PR. I also plan to keep contributing more probability-related tasks in the future.

Overview

This PR introduces a new probability module to reasoning_gym. The first dataset included is the Coin Flip dataset, which covers:

• Probability of getting exactly k heads
• Probability of getting at least k heads
• A curriculum that scales difficulty by increasing the number of coin tosses

Implementation Details

• New module: reasoning_gym/probability/coin_flip.py
  • Implements CoinFlipDataset and CoinFlipCurriculum.
• Tests: Added tests/test_coin_flip.py

Example

from reasoning_gym.probability.coin_flip import CoinFlipDataset, CoinFlipConfig

cfg = CoinFlipConfig(size=10, seed=42, min_trials=3, max_trials=15)
dataset = CoinFlipDataset(cfg)

for i in range(len(dataset)):
    ex = dataset[i]
    print("Q:", ex["question"])
    print("A:", ex["answer"])

This lays the groundwork for more probability tasks that I plan to add later (e.g., identical/non-identical objects distribution problems, expectation problems). I welcome any feedback or advice, and I’m eager to learn and improve 🙂.

Thanks, Let me know if i missed something.

[olliestanley]

Task implementation looks good I think. Left a few small comments!

[olliestanley]

Looks great, thanks!