| 933 | |
| 934 | |
| 935 | def make_replay_buffer(buffer_size, batch_size, random_crop_len, prefetch=3, prb=False): |
| 936 | if prb: |
| 937 | sampler = PrioritizedSampler( |
| 938 | max_capacity=buffer_size, |
| 939 | alpha=0.7, |
| 940 | beta=0.5, |
| 941 | ) |
| 942 | else: |
| 943 | sampler = RandomSampler() |
| 944 | replay_buffer = TensorDictReplayBuffer( |
| 945 | storage=LazyMemmapStorage( |
| 946 | buffer_size, |
| 947 | scratch_dir=buffer_scratch_dir, |
| 948 | ), |
| 949 | batch_size=batch_size, |
| 950 | sampler=sampler, |
| 951 | pin_memory=False, |
| 952 | prefetch=prefetch, |
| 953 | transform=RandomCropTensorDict(random_crop_len, sample_dim=1), |
| 954 | ) |
| 955 | return replay_buffer |
| 956 | |
| 957 | |
| 958 | ############################################################################### |