slicesampler_sandbox.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import argparse

import torch
from tensordict import TensorDict
from torch.utils.benchmark import Timer

from torchrl.data import SliceSampler, ReplayBuffer, LazyTensorStorage

parser = argparse.ArgumentParser()
parser.add_argument("--capacity", type=int, default=1_000_000)
parser.add_argument("--num_slices", type=int, default=4)
parser.add_argument("--batch_size", type=int, default=1024)

parser.add_argument("--min_traj_len", type=int, default=512)
parser.add_argument("--max_traj_len", type=int, default=1024)

parser.add_argument("--perc_full", type=float, default=0.1)
parser.add_argument("--end_or_traj", type=str, default="end")
parser.add_argument("--compile", action="store_true")
parser.add_argument("--cache_vals", action="store_true")


def main(**kwargs):
    torch.compiler.reset()

    args = parser.parse_args(**kwargs)
    if args.end_or_traj == "end":
        end_key = ("next", "done")
        traj_key = None
    else:
        traj_key = "traj_count"
        end_key = None

    if args.compile:
        compile = {"fullgraph": True}
    else:
        compile = False
    rb = ReplayBuffer(
        storage=LazyTensorStorage(args.capacity, compilable=bool(compile)),
        sampler=SliceSampler(num_slices=args.num_slices, end_key=end_key, traj_key=traj_key, compile=compile,
                             cache_values=args.cache_vals),
        batch_size=args.batch_size,
    )

    # Fill the buffer
    topval = int(args.capacity * args.perc_full)
    observation_shape = (32,)
    data = TensorDict(
        observation=torch.randn(observation_shape),
        done=torch.zeros((1,), dtype=torch.bool),
        traj_count=torch.zeros((), dtype=torch.int),
        action=torch.randint(20, ()),
        next=TensorDict(
            observation=torch.randn(observation_shape),
            done=torch.zeros((1,), dtype=torch.bool),
            reward=torch.zeros((1,)),
            traj_count=torch.zeros((), dtype=torch.int),
        )
    ).expand(topval)

    # Fill done and traj counts
    data["next", "done"] = data["next", "done"].clone()
    data["traj_count"] = data["traj_count"].clone()
    data["next", "traj_count"] = data["next", "traj_count"].clone()
    start = 0
    i = 0
    while True:
        stop = min(topval - 1, start + torch.randint(args.min_traj_len, args.max_traj_len, ()).item())
        data["next", "done"][stop] = True
        data["traj_count"][start:stop] = i
        data["next", "traj_count"][start:stop] = i
        i += 1
        if stop == topval - 1:
            break
        start = stop
    rb.extend(data)

    # warmup
    for _ in range(10):
        rb.sample()

    print(args)
    print('active storage capacity', rb[:].bytes() / 1024 / 1024 / 1024, "Gb")
    times = Timer("rb.sample()", globals={"rb": rb}).adaptive_autorange()
    print(times)
    return times.median


if __name__ == "__main__":
    import pandas as pd

    results = []

    for capacity in [10_000_000]:
        for perc_full in [0.1, 0.3, 0.5]:
            for compile in [[], ["--compile"]]:
                for cache_vals in [[], ["--cache_vals"]]:
                    median_time = main(
                        args=[f"--capacity={capacity}", f"--perc_full={perc_full}"] + compile + cache_vals)
                    results.append({
                        "Capacity": capacity,
                        "Percentage Full": perc_full,
                        "Compile": compile != [],
                        "Cache values": cache_vals != [],
                        "Median Sample Time (s)": median_time
                    })
    df = pd.DataFrame(results)
    df.set_index(["Percentage Full", "Compile", "Cache values", "Capacity"], inplace=True)
    print(df)