Change sampling method from randint to choice in Replay and robustify policy networks in SAC

rlkit/data_management/simple_replay_buffer.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict
 
 import numpy as np
+import warnings
 
 from rlkit.data_management.replay_buffer import ReplayBuffer
 
@@ -13,6 +14,7 @@ def __init__(
         observation_dim,
         action_dim,
         env_info_sizes,
+        replace = True,
     ):
         self._observation_dim = observation_dim
         self._action_dim = action_dim
@@ -35,6 +37,8 @@ def __init__(
             self._env_infos[key] = np.zeros((max_replay_buffer_size, size))
         self._env_info_keys = env_info_sizes.keys()
 
+        self._replace = replace
+
         self._top = 0
         self._size = 0
 
@@ -59,7 +63,9 @@ def _advance(self):
             self._size += 1
 
     def random_batch(self, batch_size):
-        indices = np.random.randint(0, self._size, batch_size)
+        indices = np.random.choice(self._size, size=batch_size, replace=self._replace or self._size < batch_size)
+        if not self._replace and self._size < batch_size:
+            warnings.warn('Replace was set to false, but is temporarily set to true because batch size is larger than current size of replay.')
         batch = dict(
             observations=self._observations[indices],
             actions=self._actions[indices],

rlkit/torch/sac/policies.py

@@ -111,7 +111,7 @@ def forward(
                     action,
                     pre_tanh_value=pre_tanh_value
                 )
-                log_prob = log_prob.sum(dim=1, keepdim=True)
+                log_prob = log_prob.sum(dim=-1, keepdim=True)
             else:
                 if reparameterize is True:
                     action = tanh_normal.rsample()