ray.rllib.policy.sample_batch.SampleBatch.columns

SampleBatch.columns(keys: List[str]) → List[any]

Returns a list of the batch-data in the specified columns.

Parameters:

keys – List of column names fo which to return the data.

Returns:

The list of data items ordered by the order of column names in keys.

from ray.rllib.policy.sample_batch import SampleBatch
batch = SampleBatch({"a": [1], "b": [2], "c": [3]})
print(batch.columns(["a", "b"]))

[[1], [2]]