How to simulate a contextual bandit in Python

import random


class ContextualBandit:
    def __init__(self, n_actions=3, n_features=4):
        self.n_actions = n_actions
        self.n_features = n_features
        self.theta = [random.random() for _ in range(n_actions * n_features)]

    def mock_context(self):
        return [random.uniform(-1, 1) for _ in range(self.n_features)]

    def arm_rewards(self, context):
        rewards = []
        for a in range(self.n_actions):
            theta_a = self.theta[a * self.n_features:(a + 1) * self.n_features]
            reward = sum(t * c for t, c in zip(theta_a, context))
            rewards.append(round(reward, 3))
        return rewards

    def sample_arm(self, context, epsilon=0.1):
        if random.random() < epsilon:
            return random.randint(0, self.n_actions - 1)
        rewards = self.arm_rewards(context)
        return rewards.index(max(rewards))


if __name__ == "__main__":
    bandit = ContextualBandit(n_actions=3, n_features=4)
    ctx = bandit.mock_context()
    chosen = bandit.sample_arm(ctx)
    rewards = bandit.arm_rewards(ctx)
    print(f"Context: {ctx}")
    print(f"Rewards: {rewards}")
    print(f"Chosen action (epsilon=0.1): {chosen}")