From 081fb4e6ab0a5204c0a5d1854736aa9f13f6a7a4 Mon Sep 17 00:00:00 2001
From: ebuehrle <43623224+ebuehrle@users.noreply.github.com>
Date: Thu, 4 Nov 2021 19:54:14 +0100
Subject: [PATCH] Sample with replacement

---
 scratch/etienne/intersimple/gail/options2.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/scratch/etienne/intersimple/gail/options2.py b/scratch/etienne/intersimple/gail/options2.py
index 6d3728e..69514a5 100644
--- a/scratch/etienne/intersimple/gail/options2.py
+++ b/scratch/etienne/intersimple/gail/options2.py
@@ -3,7 +3,6 @@ import torch
 from src.util.collisions import feasible
 import numpy as np
 from collections import deque
-import itertools
 
 def imitation_discriminator(discriminator):
     return lambda obs, action, next_obs, done: discriminator.discrim_net.predict_reward_train(
@@ -82,7 +81,8 @@ class OptionsEnv(gym.Wrapper):
     def sample_ll(self, n):
         assert n <= self.ll_buffer_capacity, f'Sample size of {n} exceeds buffer capacity of {self.ll_buffer_capacity}'
         assert n <= len(self.ll_buffer), f'Sample size of {n} exceeds buffer size of {len(self.ll_buffer)}'
-        return list(itertools.islice(self.ll_buffer, n))
+        ind = np.random.randint(len(self.ll_buffer), size=n)
+        return list(self.ll_buffer[i] for i in ind)
 
 class RenderOptions(OptionsEnv):