update rainbowdqn

codes/TD3/memory.py

@@ -14,13 +14,13 @@ import torch
 
 
 class ReplayBuffer(object):
-	def __init__(self, state_dim, action_dim, max_size=int(1e6)):
+	def __init__(self, n_states, n_actions, max_size=int(1e6)):
 		self.max_size = max_size
 		self.ptr = 0
 		self.size = 0
-		self.state = np.zeros((max_size, state_dim))
-		self.action = np.zeros((max_size, action_dim))
-		self.next_state = np.zeros((max_size, state_dim))
+		self.state = np.zeros((max_size, n_states))
+		self.action = np.zeros((max_size, n_actions))
+		self.next_state = np.zeros((max_size, n_states))
 		self.reward = np.zeros((max_size, 1))
 		self.not_done = np.zeros((max_size, 1))
 		self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")