Errors in documentation for Adding a new Data Bandit

[Tobias Morville]

Running the example code found here yields a shape error

from typing import Tuple

import pandas as pd
import torch

from genrl.utils.data_bandits.base import DataBasedBandit
from genrl.utils.data_bandits.utils import download_data


URL = "http://archive.ics.uci.edu/ml/machine-learning-databases/wine/wine.data"

class WineDataBandit(DataBasedBandit):
    def __init__(self, **kwargs):
        super(WineDataBandit, self).__init__(**kwargs)

        path = kwargs.get("path", "./data/Wine/")
        download = kwargs.get("download", None)
        force_download = kwargs.get("force_download", None)
        url = kwargs.get("url", URL)

        if download:
            path = download_data(path, url, force_download)

        self._df = pd.read_csv(path, header=None)
        self.n_actions = len(self._df[0].unique())
        self.context_dim = self._df.shape[1] - 1
        self.len = len(self._df)
        
        print(self.n_actions, self.context_dim, self.len)

    def reset(self) -> torch.Tensor:
        self._reset()
        self.df = self._df.sample(frac=1).reset_index(drop=True)
        return self._get_context()

    def _compute_reward(self, action: int) -> Tuple[int, int]:
        label = self._df.iloc[self.idx, 0]
        r = int(label == (action + 1))
        return r, 1

    def _get_context(self) -> torch.Tensor:
        return torch.tensor(
            self._df.iloc[self.idx, 0].values,
            device=self.device,
            dtype=torch.float,
        )

due to self._df.iloc[self.idx, 0].values returning an integer. Correcting this line to self._df.iloc[self.idx, 1:].values such that all context, except for target output, is returned makes the example work.

[Atharv Sonwane]

Thanks for pointing this out!

[Atharv Sonwane]

Yep, this is a typo, it should be self._df.iloc[self.idx, 1:].values.

Would you be up for opening a PR for the same?