def __init__(self, config, localNet, env, globalNets, globalOptimizer, netLossFunc, nbAction, rank, globalEpisodeCount, globalEpisodeReward, globalRunningAvgReward, resultQueue, logFolder, stateProcessor = None, lock = None): self.globalPolicyNet = globalNets[0] self.globalTargetNet = globalNets[1] self.rank = rank self.globalOptimizer = globalOptimizer self.localNet = localNet mp.Process.__init__(self) DQNAgent.__init__(self. config, localNet, None, None, netLossFunc, nbAction, stateProcessor, ) self.totalStep = 0 self.updateGlobalFrequency = 10 if 'updateGlobalFrequency' in self.config: self.updateGlobalFrequency = self.config['updateGlobalFrequency'] self.globalEpisodeCount = globalEpisodeCount self.globalEpisodeReward = globalEpisodeReward self.globalRunningAvgReward = globalRunningAvgReward self.resultQueue = resultQueue self.dirName = logFolder self.randomSeed = 1 + self.rank if 'randomSeed' in self.config: self.randomSeed = self.config['randomSeed'] + self.rank torch.manual_seed(self.randomSeed) self.nStepForward = 1 if 'nStepForward' in self.config: self.nStepForward = self.config['nStepForward'] self.targetNetUpdateEpisode = 10 if 'targetNetUpdateEpisode' in self.config: self.targetNetUpdateEpisode = self.config['targetNetUpdateEpisode'] self.nStepBuffer = [] # only use vanilla replay memory self.memory = ReplayMemory(self.memoryCapacity) self.priorityMemoryOption = False # use synthetic lock or not self.synchLock = False if 'synchLock' in self.config: self.synchLock = self.config['synchLock'] self.lock = lock self.device = 'cpu' if 'device' in self.config and torch.cuda.is_available(): self.device = self.config['device'] torch.cuda.manual_seed(self.randomSeed) self.localNet = self.localNet.cuda()
def init_memory(self): if self.priorityMemoryOption: self.memory = PrioritizedReplayMemory(self.memoryCapacity, self.config) else: # most commonly experience replay memory if self.memoryOption == 'natural': self.memory = ReplayMemory(self.memoryCapacity) elif self.memoryOption == 'reward': self.memory = ReplayMemoryReward(self.memoryCapacity, self.config['rewardMemoryBackupStep'], self.gamma, self.config['rewardMemoryTerminalRatio'] )
class DQNAsynERWorker(DQNAgent, mp.Process): def __init__(self, config, localNet, env, globalNets, globalOptimizer, netLossFunc, nbAction, rank, globalEpisodeCount, globalEpisodeReward, globalRunningAvgReward, resultQueue, logFolder, stateProcessor = None, lock = None): self.globalPolicyNet = globalNets[0] self.globalTargetNet = globalNets[1] self.rank = rank self.globalOptimizer = globalOptimizer self.localNet = localNet mp.Process.__init__(self) DQNAgent.__init__(self. config, localNet, None, None, netLossFunc, nbAction, stateProcessor, ) self.totalStep = 0 self.updateGlobalFrequency = 10 if 'updateGlobalFrequency' in self.config: self.updateGlobalFrequency = self.config['updateGlobalFrequency'] self.globalEpisodeCount = globalEpisodeCount self.globalEpisodeReward = globalEpisodeReward self.globalRunningAvgReward = globalRunningAvgReward self.resultQueue = resultQueue self.dirName = logFolder self.randomSeed = 1 + self.rank if 'randomSeed' in self.config: self.randomSeed = self.config['randomSeed'] + self.rank torch.manual_seed(self.randomSeed) self.nStepForward = 1 if 'nStepForward' in self.config: self.nStepForward = self.config['nStepForward'] self.targetNetUpdateEpisode = 10 if 'targetNetUpdateEpisode' in self.config: self.targetNetUpdateEpisode = self.config['targetNetUpdateEpisode'] self.nStepBuffer = [] # only use vanilla replay memory self.memory = ReplayMemory(self.memoryCapacity) self.priorityMemoryOption = False # use synthetic lock or not self.synchLock = False if 'synchLock' in self.config: self.synchLock = self.config['synchLock'] self.lock = lock self.device = 'cpu' if 'device' in self.config and torch.cuda.is_available(): self.device = self.config['device'] torch.cuda.manual_seed(self.randomSeed) self.localNet = self.localNet.cuda() def epsilon_by_episode(self, step): return self.epsilon_final + ( self.epsilon_start - self.epsilon_final) * math.exp(-1. * step / self.epsilon_decay) def run(self): torch.set_num_threads(1) bufferState, bufferAction, bufferReward, bufferNextState = [], [], [], [] for self.epIdx in range(self.trainStep): print("episode index:" + str(self.epIdx) + " from" + current_process().name + "\n") state = self.env.reset() done = False rewardSum = 0 # clear the nstep buffer self.nStepBuffer.clear() for stepCount in range(self.episodeLength): epsilon = self.epsilon_by_episode(self.globalEpisodeCount.value) action = self.select_action(self.localNet, state, epsilon) nextState, reward, done, info = self.env.step(action) if stepCount == 0: print("at step 0: from " + current_process().name + "\n") print(info) if done: nextState = None self.update_net_and_sync(state, action, nextState, reward) state = nextState rewardSum += reward * pow(self.gamma, stepCount) self.totalStep += 1 if done: # print("done in step count: {}".format(stepCount)) # print("reward sum = " + str(rewardSum)) # done and print information # pass break self.recordInfo(rewardSum, stepCount) self.resultQueue.put(None) def recordInfo(self, reward, stepCount): with self.globalEpisodeReward.get_lock(): self.globalEpisodeReward.value = reward with self.globalRunningAvgReward.get_lock(): self.globalRunningAvgReward.value = (self.globalRunningAvgReward.value * self.globalEpisodeCount.value + reward) / ( self.globalEpisodeCount.value + 1) with self.globalEpisodeCount.get_lock(): self.globalEpisodeCount.value += 1 if self.config['logFlag'] and self.globalEpisodeCount.value % self.config['logFrequency'] == 0: self.save_checkpoint() # sync global target to global policy net if self.globalEpisodeCount.value % self.targetNetUpdateEpisode == 0: self.globalTargetNet.load_state_dict(self.globalPolicyNet.state_dict()) # resultQueue.put(globalEpisodeReward.value) self.resultQueue.put( [self.globalEpisodeCount.value, stepCount, self.globalEpisodeReward.value, self.globalRunningAvgReward.value]) print(self.name) print("Episode: ", self.globalEpisodeCount.value) print("stepCount: ", stepCount) print("Episode Reward: ", self.globalEpisodeReward.value) print("Episode Running Average Reward: ", self.globalRunningAvgReward.value) def save_checkpoint(self): prefix = self.dirName + 'Epoch' + str(self.globalEpisodeCount.value) torch.save({ 'epoch': self.globalEpisodeCount.value + 1, 'model_state_dict': self.globalPolicyNet.state_dict(), 'optimizer_state_dict': self.globalOptimizer.state_dict(), }, prefix + '_checkpoint.pt') def update_net_and_sync(self, state, action, nextState, reward): self.store_experience(state, action, nextState, reward) if self.priorityMemoryOption: if len(self.memory) < self.config['memoryCapacity']: return else: if len(self.memory) < self.trainBatchSize: return if self.totalStep % self.updateGlobalFrequency == 0: transitions_raw = self.memory.sample(self.trainBatchSize) transitions = Transition(*zip(*transitions_raw)) action = torch.tensor(transitions.action, device=self.device, dtype=torch.long).unsqueeze( -1) # shape(batch, 1) reward = torch.tensor(transitions.reward, device=self.device, dtype=torch.float32).unsqueeze( -1) # shape(batch, 1) batchSize = reward.shape[0] # for some env, the output state requires further processing before feeding to neural network if self.stateProcessor is not None: state, _ = self.stateProcessor(transitions.state, self.device) nonFinalNextState, nonFinalMask = self.stateProcessor(transitions.next_state, self.device) else: state = torch.tensor(transitions.state, device=self.device, dtype=torch.float32) nonFinalMask = torch.tensor([s is not None for s in transitions.next_state], device=self.device, dtype=torch.uint8) nonFinalNextState = torch.tensor([s for s in transitions.next_state if s is not None], device=self.device, dtype=torch.float32) if self.synchLock: self.lock.acquire() QValues = self.globalPolicyNet(state).gather(1, action) if self.netUpdateOption == 'targetNet': # Here we detach because we do not want gradient flow from target values to net parameters QNext = torch.zeros(batchSize, device=self.device, dtype=torch.float32) QNext[nonFinalMask] = self.globalTargetNet(nonFinalNextState).max(1)[0].detach() targetValues = reward + self.gamma * QNext.unsqueeze(-1) if self.netUpdateOption == 'policyNet': raise NotImplementedError targetValues = reward + self.gamma * torch.max(self.globalPolicyNet(nextState).detach(), dim=1)[0].unsqueeze(-1) if self.netUpdateOption == 'doubleQ': # select optimal action from policy net with torch.no_grad(): batchAction = self.globalPolicyNet(nonFinalNextState).max(dim=1)[1].unsqueeze(-1) QNext = torch.zeros(batchSize, device=self.device, dtype=torch.float32).unsqueeze(-1) QNext[nonFinalMask] = self.globalTargetNet(nonFinalNextState).gather(1, batchAction) targetValues = reward + self.gamma * QNext loss = self.netLossFunc(QValues, targetValues) self.globalOptimizer.zero_grad() loss.backward() if self.netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.globalPolicyNet.parameters(), self.netGradClip) # global net update self.globalOptimizer.step() # # # update local net self.localNet.load_state_dict(self.globalPolicyNet.state_dict()) self.lock.release() else: # update local net self.localNet.load_state_dict(self.globalPolicyNet.state_dict()) QValues = self.localNet(state).gather(1, action) if self.netUpdateOption == 'targetNet': # Here we detach because we do not want gradient flow from target values to net parameters QNext = torch.zeros(batchSize, device=self.device, dtype=torch.float32) QNext[nonFinalMask] = self.globalTargetNet(nonFinalNextState).max(1)[0].detach() targetValues = reward + self.gamma * QNext.unsqueeze(-1) if self.netUpdateOption == 'policyNet': raise NotImplementedError targetValues = reward + self.gamma * torch.max(self.globalPolicyNet(nextState).detach(), dim=1)[ 0].unsqueeze(-1) if self.netUpdateOption == 'doubleQ': # select optimal action from policy net with torch.no_grad(): batchAction = self.localNet(nonFinalNextState).max(dim=1)[1].unsqueeze(-1) QNext = torch.zeros(batchSize, device=self.device, dtype=torch.float32).unsqueeze(-1) QNext[nonFinalMask] = self.globalTargetNet(nonFinalNextState).gather(1, batchAction) targetValues = reward + self.gamma * QNext loss = self.netLossFunc(QValues, targetValues) loss.backward() self.lock.acquire() self.globalOptimizer.zero_grad() for lp, gp in zip(self.localNet.parameters(), self.globalPolicyNet.parameters()): if self.device == 'cpu': gp._grad = lp._grad else: gp._grad = lp._grad.cpu() if self.netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.globalPolicyNet.parameters(), self.netGradClip) # global net update self.globalOptimizer.step() self.lock.release() # # # update local net self.localNet.load_state_dict(self.globalPolicyNet.state_dict()) def test_multiProcess(self): print("Hello, World! from " + current_process().name + "\n") print(self.globalPolicyNet.state_dict()) for gp in self.globalPolicyNet.parameters(): gp.grad = torch.ones_like(gp) #gp.grad.fill_(1) self.globalOptimizer.step() print('globalNetID:') print(id(self.globalPolicyNet)) print('globalOptimizer:') print(id(self.globalOptimizer)) print('localNetID:') print(id(self.localNet))
def init_memory(self): self.memory = ReplayMemory(self.memoryCapacity)
class TDDDPGAgent(DDPGAgent): def __init__(self, config, actorNets, criticNets, env, optimizers, netLossFunc, nbAction, stateProcessor=None, experienceProcessor=None): super(TDDDPGAgent, self).__init__(config, actorNets, criticNets, env, optimizers, netLossFunc, nbAction, stateProcessor, experienceProcessor) def initalizeNets(self, actorNets, criticNets, optimizers): self.actorNet = actorNets['actor'] self.actorNet_target = actorNets[ 'target'] if 'target' in actorNets else None self.criticNetOne = criticNets['criticOne'] self.criticNet_targetOne = criticNets[ 'targetOne'] if 'targetOne' in criticNets else None self.criticNetTwo = criticNets['criticTwo'] self.criticNet_targetTwo = criticNets[ 'targetTwo'] if 'targetTwo' in criticNets else None self.actor_optimizer = optimizers['actor'] self.criticOne_optimizer = optimizers['criticOne'] self.criticTwo_optimizer = optimizers['criticTwo'] self.net_to_device() def init_memory(self): self.memory = ReplayMemory(self.memoryCapacity) def read_config(self): super(TDDDPGAgent, self).read_config() self.policyUpdateFreq = 2 if 'policyUpdateFreq' in self.config: self.policyUpdateFreq = self.config['policyUpdateFreq'] self.policySmoothNoise = 0.01 if 'policySmoothNoise' in self.config: self.policyUpdateFreq = self.config['policySmoothNoise'] def net_to_device(self): # move model to correct device self.actorNet = self.actorNet.to(self.device) self.criticNetOne = self.criticNetOne.to(self.device) self.criticNetTwo = self.criticNetTwo.to(self.device) # in case targetNet is None if self.actorNet_target is not None: self.actorNet_target = self.actorNet_target.to(self.device) # in case targetNet is None if self.criticNet_targetOne is not None: self.criticNet_targetOne = self.criticNet_targetOne.to(self.device) if self.criticNet_targetTwo is not None: self.criticNet_targetTwo = self.criticNet_targetTwo.to(self.device) def prepare_minibatch(self, state, action, nextState, reward, info): # first store memory self.store_experience(state, action, nextState, reward, info) if len(self.memory) < self.trainBatchSize: return transitions_raw = self.memory.sample(self.trainBatchSize) transitions = Transition(*zip(*transitions_raw)) action = torch.tensor(transitions.action, device=self.device, dtype=torch.float32) # shape(batch, numActions) reward = torch.tensor(transitions.reward, device=self.device, dtype=torch.float32) # shape(batch) # for some env, the output state requires further processing before feeding to neural network if self.stateProcessor is not None: state, _ = self.stateProcessor(transitions.state, self.device) nonFinalNextState, nonFinalMask = self.stateProcessor( transitions.next_state, self.device) else: state = torch.tensor(transitions.state, device=self.device, dtype=torch.float32) nonFinalMask = torch.tensor(tuple( map(lambda s: s is not None, transitions.next_state)), device=self.device, dtype=torch.uint8) nonFinalNextState = torch.tensor( [s for s in transitions.next_state if s is not None], device=self.device, dtype=torch.float32) return state, nonFinalMask, nonFinalNextState, action, reward def update_net(self, state, action, nextState, reward, info): # state, nonFinalMask, nonFinalNextState, action, reward = self.prepare_minibatch(state, action, nextState, reward, info) self.store_experience(state, action, nextState, reward, info) if len(self.memory) < self.trainBatchSize: return transitions_raw = self.memory.sample(self.trainBatchSize) transitions = Transition(*zip(*transitions_raw)) action = torch.tensor(transitions.action, device=self.device, dtype=torch.float32) # shape(batch, numActions) reward = torch.tensor(transitions.reward, device=self.device, dtype=torch.float32) # shape(batch) # for some env, the output state requires further processing before feeding to neural network if self.stateProcessor is not None: state, _ = self.stateProcessor(transitions.state, self.device) nonFinalNextState, nonFinalMask = self.stateProcessor( transitions.next_state, self.device) else: state = torch.tensor(transitions.state, device=self.device, dtype=torch.float32) nonFinalMask = torch.tensor(tuple( map(lambda s: s is not None, transitions.next_state)), device=self.device, dtype=torch.uint8) nonFinalNextState = torch.tensor( [s for s in transitions.next_state if s is not None], device=self.device, dtype=torch.float32) batchSize = reward.shape[0] # Critic loss QValuesOne = self.criticNetOne.forward(state, action).squeeze() QValuesTwo = self.criticNetTwo.forward(state, action).squeeze() actionNoise = torch.randn((nonFinalNextState.shape[0], self.numAction), dtype=torch.float32, device=self.device) next_actions = self.actorNet_target.forward( nonFinalNextState) + actionNoise * self.policySmoothNoise # next_actions = self.actorNet_target.forward(nonFinalNextState) QNext = torch.zeros(batchSize, device=self.device, dtype=torch.float32) QNextCriticOne = self.criticNet_targetOne.forward( nonFinalNextState, next_actions.detach()).squeeze() QNextCriticTwo = self.criticNet_targetTwo.forward( nonFinalNextState, next_actions.detach()).squeeze() QNext[nonFinalMask] = torch.min(QNextCriticOne, QNextCriticTwo) targetValues = reward + self.gamma * QNext criticOne_loss = self.netLossFunc(QValuesOne, targetValues) criticTwo_loss = self.netLossFunc(QValuesTwo, targetValues) self.criticOne_optimizer.zero_grad() self.criticTwo_optimizer.zero_grad() # https://jdhao.github.io/2017/11/12/pytorch-computation-graph/ criticOne_loss.backward(retain_graph=True) criticTwo_loss.backward() if self.netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.criticNetOne.parameters(), self.netGradClip) torch.nn.utils.clip_grad_norm_(self.criticNetTwo.parameters(), self.netGradClip) self.criticOne_optimizer.step() self.criticTwo_optimizer.step() if self.learnStepCounter % self.policyUpdateFreq: # Actor loss # we try to maximize criticNet output(which is state value) policy_loss = -self.criticNetOne.forward( state, self.actorNet.forward(state)).mean() # update networks self.actor_optimizer.zero_grad() policy_loss.backward() if self.netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.actorNet.parameters(), self.netGradClip) self.actor_optimizer.step() if self.globalStepCount % self.lossRecordStep == 0: self.losses.append([ self.globalStepCount, self.epIdx, criticOne_loss.item(), criticTwo_loss.item(), policy_loss.item() ]) # update target networks for target_param, param in zip( self.actorNet_target.parameters(), self.actorNet.parameters()): target_param.data.copy_(param.data * self.tau + target_param.data * (1.0 - self.tau)) for target_param, param in zip( self.criticNet_targetOne.parameters(), self.criticNetOne.parameters()): target_param.data.copy_(param.data * self.tau + target_param.data * (1.0 - self.tau)) for target_param, param in zip( self.criticNet_targetTwo.parameters(), self.criticNetTwo.parameters()): target_param.data.copy_(param.data * self.tau + target_param.data * (1.0 - self.tau)) self.learnStepCounter += 1 def save_all(self): prefix = self.dirName + self.identifier + 'Finalepoch' + str( self.epIdx) self.saveLosses(prefix + '_loss.txt') self.saveRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'wb') as file: pickle.dump(self.memory, file) torch.save( { 'epoch': self.epIdx, 'globalStep': self.globalStepCount, 'actorNet_state_dict': self.actorNet.state_dict(), 'criticNetOne_state_dict': self.criticNetOne.state_dict(), 'criticNetTwo_state_dict': self.criticNetTwo.state_dict(), 'actor_optimizer_state_dict': self.actor_optimizer.state_dict(), 'criticOne_optimizer_state_dict': self.criticOne_optimizer.state_dict(), 'criticTwo_optimizer_state_dict': self.criticOne_optimizer.state_dict() }, prefix + '_checkpoint.pt') def save_checkpoint(self): prefix = self.dirName + self.identifier + 'Epoch' + str(self.epIdx) self.saveLosses(prefix + '_loss.txt') self.saveRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'wb') as file: pickle.dump(self.memory, file) torch.save( { 'epoch': self.epIdx, 'globalStep': self.globalStepCount, 'actorNet_state_dict': self.actorNet.state_dict(), 'criticNetOne_state_dict': self.criticNetOne.state_dict(), 'criticNetTwo_state_dict': self.criticNetTwo.state_dict(), 'actor_optimizer_state_dict': self.actor_optimizer.state_dict(), 'criticOne_optimizer_state_dict': self.criticOne_optimizer.state_dict(), 'criticTwo_optimizer_state_dict': self.criticTwo_optimizer.state_dict() }, prefix + '_checkpoint.pt') def load_checkpoint(self, prefix): self.loadLosses(prefix + '_loss.txt') self.loadRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'rb') as file: self.memory = pickle.load(file) checkpoint = torch.load(prefix + '_checkpoint.pt') self.epIdx = checkpoint['epoch'] self.globalStepCount = checkpoint['globalStep'] self.actorNet.load_state_dict(checkpoint['actorNet_state_dict']) self.actorNet_target.load_state_dict(checkpoint['actorNet_state_dict']) self.criticNetOne.load_state_dict( checkpoint['criticNetOne_state_dict']) self.criticNet_targetOne.load_state_dict( checkpoint['criticNetOne_state_dict']) self.criticNetTwo.load_state_dict( checkpoint['criticNetTwo_state_dict']) self.criticNet_targetTwo.load_state_dict( checkpoint['criticNetTwo_state_dict']) self.actor_optimizer.load_state_dict( checkpoint['actor_optimizer_state_dict']) self.criticOne_optimizer.load_state_dict( checkpoint['criticOne_optimizer_state_dict']) self.criticTwo_optimizer.load_state_dict( checkpoint['criticTwo_optimizer_state_dict'])
def init_memory(self): self.memories = [ReplayMemory(self.memoryCapacity) for _ in range(self.episodeLength)]
from Agents.Core.ReplayMemory import ReplayMemory, Transition #from ..Agents.Core.ReplayMemory import ReplayMemory, Transition import torch import numpy as np import pickle state1 = np.random.rand(5, 5) state2 = np.random.rand(5, 5) state3 = np.random.rand(5, 5) state4 = np.random.rand(5, 5) tran1 = Transition(state1, 1, state2, 1) tran2 = Transition(state3, 2, state4, 2) memory = ReplayMemory(10) memory.push(tran1) memory.push(tran2) print(memory) file = open('memory.pickle', 'wb') pickle.dump(memory, file) file.close() with open('memory.pickle', 'rb') as file: memory2 = pickle.load(file) print(memory2)
from Agents.Core.ReplayMemory import ReplayMemory, Transition #from ..Agents.Core.ReplayMemory import ReplayMemory, Transition import torch tran1 = Transition(1, 1, 1, 1) tran2 = Transition(2, 2, 2, 2) memory = ReplayMemory(10) memory.push(tran1) memory.push(tran2) memory.push(3, 3, 3, 3) print(memory) memory.write_to_text('memoryOut.txt') toTensor = memory.totensor() toTensor2 = torch.tensor(memory.sample(2)) for i in range(5, 50): tran = Transition(i, i, i, i) memory.push(tran) print(memory) memory.clear() print(memory) print(toTensor) print(toTensor2)
env, optimizer, torch.nn.MSELoss(reduction='none'), N_A, config=config) xSet = np.linspace(-1, 1, 100) policy = np.zeros_like(xSet) for i, x in enumerate(xSet): policy[i] = agent.getPolicy(np.array([x])) np.savetxt('StabilizerPolicyBeforeTrain.txt', policy, fmt='%d') #agent.perform_random_exploration(10) agent.train() storeMemory = ReplayMemory(100000) agent.testPolicyNet(100, storeMemory) storeMemory.write_to_text('testPolicyMemory.txt') def customPolicy(state): x = state[0] # move towards negative if x > 0.1: action = 2 # move towards positive elif x < -0.1: action = 1 # do not move else: action = 0
class DDPGAgent: """class for DDPG agents. This class contains implementation of DDPG learning. It contains enhancement of experience augmentation, hindsight experience replay. # Arguments config: a dictionary for training parameters actors: actor net and its target net criticNets: critic net and its target net env: environment for the agent to interact. env should implement same interface of a gym env optimizers: network optimizers for both actor net and critic netLossFunc: loss function of the network, e.g., mse nbAction: number of actions stateProcessor: a function to process output from env, processed state will be used as input to the networks experienceProcessor: additional steps to process an experience """ def __init__(self, config, actorNets, criticNets, env, optimizers, netLossFunc, nbAction, stateProcessor=None, experienceProcessor=None): self.config = config self.read_config() self.env = env self.numAction = nbAction self.stateProcessor = stateProcessor self.netLossFunc = netLossFunc self.experienceProcessor = experienceProcessor self.initialization() self.init_memory() self.initalizeNets(actorNets, criticNets, optimizers) def initalizeNets(self, actorNets, criticNets, optimizers): ''' initialize networks and their optimizers; move them to specified device (i.e., cpu or cuda) ''' self.actorNet = actorNets['actor'] self.actorNet_target = actorNets[ 'target'] if 'target' in actorNets else None self.criticNet = criticNets['critic'] self.criticNet_target = criticNets[ 'target'] if 'target' in criticNets else None self.actor_optimizer = optimizers['actor'] self.critic_optimizer = optimizers['critic'] self.net_to_device() def init_memory(self): ''' initialize replay memory ''' self.memory = ReplayMemory(self.memoryCapacity) def read_config(self): ''' read parameters from self.config object initialize various flags and parameters trainStep: number of episodes to train targetNetUpdateStep: frequency in terms of training steps/episodes to reset target net trainBatchSize: mini batch size for gradient decent. gamma: discount factor tau: soft update parameter memoryCapacity: memory capacity for experience storage netGradClip: gradient clipping parameter netUpdateOption: allowed strings are targetNet, policyNet, doubleQ verbose: bool, default false. nStepForward: multiple-step forward Q learning, default 1 lossRecordStep: frequency to store loss. episodeLength: maximum steps in an episode netUpdateFrequency: frequency to perform gradient decent netUpdateStep: number of steps for gradient decent device: cpu or cuda randomSeed hindSightER: bool variable for hindsight experience replay hindSightERFreq: frequency to perform hindsight experience replay experienceAugmentation: additional experience augmentation function return: None ''' self.trainStep = self.config['trainStep'] self.targetNetUpdateStep = 10000 if 'targetNetUpdateStep' in self.config: self.targetNetUpdateStep = self.config['targetNetUpdateStep'] self.trainBatchSize = self.config['trainBatchSize'] self.gamma = self.config['gamma'] self.tau = self.config['tau'] self.netGradClip = None if 'netGradClip' in self.config: self.netGradClip = self.config['netGradClip'] self.netUpdateOption = 'targetNet' if 'netUpdateOption' in self.config: self.netUpdateOption = self.config['netUpdateOption'] self.verbose = False if 'verbose' in self.config: self.verbose = self.config['verbose'] self.netUpdateFrequency = 1 if 'netUpdateFrequency' in self.config: self.netUpdateFrequency = self.config['netUpdateFrequency'] self.nStepForward = 1 if 'nStepForward' in self.config: self.nStepForward = self.config['nStepForward'] self.lossRecordStep = 10 if 'lossRecordStep' in self.config: self.lossRecordStep = self.config['lossRecordStep'] self.episodeLength = 500 if 'episodeLength' in self.config: self.episodeLength = self.config['episodeLength'] self.verbose = False if 'verbose' in self.config: self.verbose = self.config['verbose'] self.device = 'cpu' if 'device' in self.config and torch.cuda.is_available(): self.device = self.config['device'] self.randomSeed = 1 if 'randomSeed' in self.config: self.randomSeed = self.config['randomSeed'] self.memoryCapacity = self.config['memoryCapacity'] self.hindSightER = False if 'hindSightER' in self.config: self.hindSightER = self.config['hindSightER'] self.hindSightERFreq = self.config['hindSightERFreq'] self.experienceAugmentation = False if 'experienceAugmentation' in self.config: self.experienceAugmentation = self.config['experienceAugmentation'] self.experienceAugmentationFreq = self.config[ 'experienceAugmentationFreq'] self.policyUpdateFreq = 1 if 'policyUpdateFreq' in self.config: self.policyUpdateFreq = self.config['policyUpdateFreq'] def net_to_device(self): ''' move model to the specified devices ''' self.actorNet = self.actorNet.to(self.device) self.criticNet = self.criticNet.to(self.device) # in case targetNet is None if self.actorNet_target is not None: self.actorNet_target = self.actorNet_target.to(self.device) # in case targetNet is None if self.criticNet_target is not None: self.criticNet_target = self.criticNet_target.to(self.device) def initialization(self): self.dirName = 'Log/' if 'dataLogFolder' in self.config: self.dirName = self.config['dataLogFolder'] if not os.path.exists(self.dirName): os.makedirs(self.dirName) self.identifier = '' self.epIdx = 0 self.learnStepCounter = 0 # for target net update self.globalStepCount = 0 self.losses = [] self.rewards = [] self.runningAvgEpisodeReward = 0.0 def select_action(self, net=None, state=None, noiseFlag=False): ''' select action from net. The action selection is delegated to network to implement the method of 'select_action' # Arguments net: which net used for action selection. default is actorNet state: observation or state as the input to the net noiseFlag: if set False, will perform greedy selection. if True, will add noise from OU processes. return: numpy array of actions ''' if net is None: net = self.actorNet with torch.no_grad(): # self.policyNet(torch.from_numpy(state.astype(np.float32)).unsqueeze(0)) # here state[np.newaxis,:] is to add a batch dimension if self.stateProcessor is not None: state, _ = self.stateProcessor([state], self.device) action = net.select_action(state, noiseFlag) else: stateTorch = torch.from_numpy( np.array(state[np.newaxis, :], dtype=np.float32)) action = net.select_action(stateTorch.to(self.device), noiseFlag) return action.cpu().data.numpy()[0] def process_hindSightExperience(self, state, action, nextState, reward, info): if nextState is not None and self.globalStepCount % self.hindSightERFreq == 0: stateNew, actionNew, nextStateNew, rewardNew = self.env.getHindSightExperience( state, action, nextState, info) if stateNew is not None: transition = Transition(stateNew, actionNew, nextStateNew, rewardNew) self.memory.push(transition) if self.experienceAugmentation: self.process_experienceAugmentation( state, action, nextState, reward, info) def process_experienceAugmentation(self, state, action, nextState, reward, info): if self.globalStepCount % self.experienceAugmentationFreq == 0: state_Augs, action_Augs, nextState_Augs, reward_Augs = self.env.getExperienceAugmentation( state, action, nextState, reward, info) for i in range(len(state_Augs)): transition = Transition(state_Augs[i], action_Augs[i], nextState_Augs[i], reward_Augs[i]) self.memory.push(transition) def store_experience(self, state, action, nextState, reward, info): if self.experienceProcessor is not None: state, action, nextState, reward = self.experienceProcessor( state, action, nextState, reward, info) transition = Transition(state, action, nextState, reward) self.memory.push(transition) if self.experienceAugmentation: self.process_experienceAugmentation(state, action, nextState, reward, info) if self.hindSightER: self.process_hindSightExperience(state, action, nextState, reward, info) def prepare_minibatch(self, transitions_raw): ''' do some proprocessing work for transitions_raw order the data convert transition list to torch tensors use trick from https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html https://stackoverflow.com/questions/19339/transpose-unzip-function-inverse-of-zip/19343#19343 ''' transitions = Transition(*zip(*transitions_raw)) action = torch.tensor(transitions.action, device=self.device, dtype=torch.float32) # shape(batch, numActions) reward = torch.tensor(transitions.reward, device=self.device, dtype=torch.float32) # shape(batch) # for some env, the output state requires further processing before feeding to neural network if self.stateProcessor is not None: state, _ = self.stateProcessor(transitions.state, self.device) nonFinalNextState, nonFinalMask = self.stateProcessor( transitions.next_state, self.device) else: state = torch.tensor(transitions.state, device=self.device, dtype=torch.float32) nonFinalMask = torch.tensor(tuple( map(lambda s: s is not None, transitions.next_state)), device=self.device, dtype=torch.bool) nonFinalNextState = torch.tensor( [s for s in transitions.next_state if s is not None], device=self.device, dtype=torch.float32) return state, nonFinalMask, nonFinalNextState, action, reward def update_net(self, state, action, nextState, reward, info): ''' This routine will store, transform, augment experiences and sample experiences for gradient descent. ''' self.store_experience(state, action, nextState, reward, info) # prepare mini-batch if len(self.memory) < self.trainBatchSize: return transitions_raw = self.memory.sample(self.trainBatchSize) self.update_net_on_transitions(transitions_raw) self.copy_nets() self.learnStepCounter += 1 def copy_nets(self): ''' soft update target networks ''' # update networks if self.learnStepCounter % self.policyUpdateFreq == 0: # update target networks for target_param, param in zip(self.actorNet_target.parameters(), self.actorNet.parameters()): target_param.data.copy_(param.data * self.tau + target_param.data * (1.0 - self.tau)) for target_param, param in zip(self.criticNet_target.parameters(), self.criticNet.parameters()): target_param.data.copy_(param.data * self.tau + target_param.data * (1.0 - self.tau)) def update_net_on_transitions(self, transitions_raw): ''' This function performs gradient gradient on the network ''' state, nonFinalMask, nonFinalNextState, action, reward = self.prepare_minibatch( transitions_raw) # Critic loss QValues = self.criticNet.forward(state, action).squeeze() QNext = torch.zeros(self.trainBatchSize, device=self.device, dtype=torch.float32) if len(nonFinalNextState): # next action is calculated using target actor network next_actions = self.actorNet_target.forward(nonFinalNextState) QNext[nonFinalMask] = self.criticNet_target.forward( nonFinalNextState, next_actions.detach()).squeeze() targetValues = reward + self.gamma * QNext critic_loss = self.netLossFunc(QValues, targetValues) self.critic_optimizer.zero_grad() critic_loss.backward() if self.netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.criticNet.parameters(), self.netGradClip) self.critic_optimizer.step() # update networks if self.learnStepCounter % self.policyUpdateFreq == 0: # Actor loss # we try to maximize criticNet output(which is state value) policy_loss = -self.criticNet.forward( state, self.actorNet.forward(state)).mean() self.actor_optimizer.zero_grad() policy_loss.backward() if self.netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.actorNet.parameters(), self.netGradClip) self.actor_optimizer.step() if self.globalStepCount % self.lossRecordStep == 0: self.losses.append([ self.globalStepCount, self.epIdx, critic_loss.item(), policy_loss.item() ]) def work_before_step(self, state=None): pass def train_one_episode(self): print("episode index:" + str(self.epIdx)) state = self.env.reset() done = False rewardSum = 0 for stepCount in range(self.episodeLength): # any work to be done before select actions self.work_before_step(state) action = self.select_action(self.actorNet, state, noiseFlag=True) nextState, reward, done, info = self.env.step(action) if stepCount == 0: print("at step 0:") print(info) if done: nextState = None # learn the transition self.update_net(state, action, nextState, reward, info) state = nextState rewardSum += reward * pow(self.gamma, stepCount) self.globalStepCount += 1 if self.verbose: print('action: ' + str(action)) print('state:') print(nextState) print('reward:') print(reward) print('info') print(info) if done: break self.runningAvgEpisodeReward = ( self.runningAvgEpisodeReward * self.epIdx + rewardSum) / (self.epIdx + 1) print("done in step count: {}".format(stepCount)) print("reward sum = " + str(rewardSum)) print("running average episode reward sum: {}".format( self.runningAvgEpisodeReward)) print(info) self.rewards.append([ self.epIdx, stepCount, self.globalStepCount, rewardSum, self.runningAvgEpisodeReward ]) if self.config[ 'logFlag'] and self.epIdx % self.config['logFrequency'] == 0: self.save_checkpoint() self.epIdx += 1 def train(self): # continue on historical training if len(self.rewards) > 0: self.runningAvgEpisodeReward = self.rewards[-1][-1] for trainStepCount in range(self.trainStep): self.train_one_episode() self.save_all() def saveLosses(self, fileName): np.savetxt(fileName, np.array(self.losses), fmt='%.5f', delimiter='\t') def saveRewards(self, fileName): np.savetxt(fileName, np.array(self.rewards), fmt='%.5f', delimiter='\t') def loadLosses(self, fileName): self.losses = np.genfromtxt(fileName).tolist() def loadRewards(self, fileName): self.rewards = np.genfromtxt(fileName).tolist() def save_all(self, identifier=None): if identifier is None: identifier = self.identifier prefix = self.dirName + identifier + 'Finalepoch' + str(self.epIdx) self.saveLosses(prefix + '_loss.txt') self.saveRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'wb') as file: pickle.dump(self.memory, file) torch.save( { 'epoch': self.epIdx, 'globalStep': self.globalStepCount, 'actorNet_state_dict': self.actorNet.state_dict(), 'criticNet_state_dict': self.criticNet.state_dict(), 'actor_optimizer_state_dict': self.actor_optimizer.state_dict(), 'critic_optimizer_state_dict': self.critic_optimizer.state_dict() }, prefix + '_checkpoint.pt') def save_checkpoint(self, identifier=None): if identifier is None: identifier = self.identifier prefix = self.dirName + identifier + 'Epoch' + str(self.epIdx) self.saveLosses(prefix + '_loss.txt') self.saveRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'wb') as file: pickle.dump(self.memory, file) torch.save( { 'epoch': self.epIdx, 'globalStep': self.globalStepCount, 'actorNet_state_dict': self.actorNet.state_dict(), 'criticNet_state_dict': self.criticNet.state_dict(), 'actor_optimizer_state_dict': self.actor_optimizer.state_dict(), 'critic_optimizer_state_dict': self.critic_optimizer.state_dict() }, prefix + '_checkpoint.pt') def load_checkpoint(self, prefix): self.loadLosses(prefix + '_loss.txt') self.loadRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'rb') as file: self.memory = pickle.load(file) checkpoint = torch.load(prefix + '_checkpoint.pt') self.epIdx = checkpoint['epoch'] self.globalStepCount = checkpoint['globalStep'] self.actorNet.load_state_dict(checkpoint['actorNet_state_dict']) self.actorNet_target.load_state_dict(checkpoint['actorNet_state_dict']) self.criticNet.load_state_dict(checkpoint['criticNet_state_dict']) self.criticNet_target.load_state_dict( checkpoint['criticNet_state_dict']) self.actor_optimizer.load_state_dict( checkpoint['actor_optimizer_state_dict']) self.critic_optimizer.load_state_dict( checkpoint['critic_optimizer_state_dict'])
def init_memory(self): ''' initialize replay memory ''' self.memory = ReplayMemory(self.memoryCapacity)
class DDPGAgent: def __init__(self, config, actorNets, criticNets, env, optimizers, netLossFunc, nbAction, stateProcessor=None, experienceProcessor=None): self.config = config self.read_config() self.env = env self.numAction = nbAction self.stateProcessor = stateProcessor self.netLossFunc = netLossFunc self.experienceProcessor = experienceProcessor self.initialization() self.init_memory() self.initalizeNets(actorNets, criticNets, optimizers) def initalizeNets(self, actorNets, criticNets, optimizers): self.actorNet = actorNets['actor'] self.actorNet_target = actorNets[ 'target'] if 'target' in actorNets else None self.criticNet = criticNets['critic'] self.criticNet_target = criticNets[ 'target'] if 'target' in criticNets else None self.actor_optimizer = optimizers['actor'] self.critic_optimizer = optimizers['critic'] self.net_to_device() def init_memory(self): self.memory = ReplayMemory(self.memoryCapacity) def read_config(self): self.trainStep = self.config['trainStep'] self.targetNetUpdateStep = 10000 if 'targetNetUpdateStep' in self.config: self.targetNetUpdateStep = self.config['targetNetUpdateStep'] self.trainBatchSize = self.config['trainBatchSize'] self.gamma = self.config['gamma'] self.tau = self.config['tau'] self.netGradClip = None if 'netGradClip' in self.config: self.netGradClip = self.config['netGradClip'] self.netUpdateOption = 'targetNet' if 'netUpdateOption' in self.config: self.netUpdateOption = self.config['netUpdateOption'] self.verbose = False if 'verbose' in self.config: self.verbose = self.config['verbose'] self.netUpdateFrequency = 1 if 'netUpdateFrequency' in self.config: self.netUpdateFrequency = self.config['netUpdateFrequency'] self.nStepForward = 1 if 'nStepForward' in self.config: self.nStepForward = self.config['nStepForward'] self.lossRecordStep = 10 if 'lossRecordStep' in self.config: self.lossRecordStep = self.config['lossRecordStep'] self.episodeLength = 500 if 'episodeLength' in self.config: self.episodeLength = self.config['episodeLength'] self.verbose = False if 'verbose' in self.config: self.verbose = self.config['verbose'] self.device = 'cpu' if 'device' in self.config and torch.cuda.is_available(): self.device = self.config['device'] self.randomSeed = 1 if 'randomSeed' in self.config: self.randomSeed = self.config['randomSeed'] self.memoryCapacity = self.config['memoryCapacity'] self.hindSightER = False if 'hindSightER' in self.config: self.hindSightER = self.config['hindSightER'] self.hindSightERFreq = self.config['hindSightERFreq'] self.experienceAugmentation = False if 'experienceAugmentation' in self.config: self.experienceAugmentation = self.config['experienceAugmentation'] self.experienceAugmentationFreq = self.config[ 'experienceAugmentationFreq'] self.policyUpdateFreq = 1 if 'policyUpdateFreq' in self.config: self.policyUpdateFreq = self.config['policyUpdateFreq'] def net_to_device(self): # move model to correct device self.actorNet = self.actorNet.to(self.device) self.criticNet = self.criticNet.to(self.device) # in case targetNet is None if self.actorNet_target is not None: self.actorNet_target = self.actorNet_target.to(self.device) # in case targetNet is None if self.criticNet_target is not None: self.criticNet_target = self.criticNet_target.to(self.device) def initialization(self): self.dirName = 'Log/' if 'dataLogFolder' in self.config: self.dirName = self.config['dataLogFolder'] if not os.path.exists(self.dirName): os.makedirs(self.dirName) self.identifier = '' self.epIdx = 0 self.learnStepCounter = 0 # for target net update self.globalStepCount = 0 self.losses = [] self.rewards = [] def select_action(self, net, state, noiseFlag=False): with torch.no_grad(): # self.policyNet(torch.from_numpy(state.astype(np.float32)).unsqueeze(0)) # here state[np.newaxis,:] is to add a batch dimension if self.stateProcessor is not None: state, _ = self.stateProcessor([state], self.device) action = net.select_action(state, noiseFlag) else: stateTorch = torch.from_numpy( np.array(state[np.newaxis, :], dtype=np.float32)) action = net.select_action(stateTorch.to(self.device), noiseFlag) return action.cpu().data.numpy()[0] def process_hindSightExperience(self, state, action, nextState, reward, info): if nextState is not None and self.globalStepCount % self.hindSightERFreq == 0: stateNew, actionNew, nextStateNew, rewardNew = self.env.getHindSightExperience( state, action, nextState, info) if stateNew is not None: transition = Transition(stateNew, actionNew, nextStateNew, rewardNew) self.memory.push(transition) if self.experienceAugmentation: self.process_experienceAugmentation( state, action, nextState, reward, info) def process_experienceAugmentation(self, state, action, nextState, reward, info): if self.globalStepCount % self.experienceAugmentationFreq == 0: state_Augs, action_Augs, nextState_Augs, reward_Augs = self.env.getExperienceAugmentation( state, action, nextState, reward, info) for i in range(len(state_Augs)): transition = Transition(state_Augs[i], action_Augs[i], nextState_Augs[i], reward_Augs[i]) self.memory.push(transition) def store_experience(self, state, action, nextState, reward, info): if self.experienceProcessor is not None: state, action, nextState, reward = self.experienceProcessor( state, action, nextState, reward, info) transition = Transition(state, action, nextState, reward) self.memory.push(transition) if self.experienceAugmentation: self.process_experienceAugmentation(state, action, nextState, reward, info) if self.hindSightER: self.process_hindSightExperience(state, action, nextState, reward, info) def prepare_minibatch(self, transitions_raw): # first store memory transitions = Transition(*zip(*transitions_raw)) action = torch.tensor(transitions.action, device=self.device, dtype=torch.float32) # shape(batch, numActions) reward = torch.tensor(transitions.reward, device=self.device, dtype=torch.float32) # shape(batch) # for some env, the output state requires further processing before feeding to neural network if self.stateProcessor is not None: state, _ = self.stateProcessor(transitions.state, self.device) nonFinalNextState, nonFinalMask = self.stateProcessor( transitions.next_state, self.device) else: state = torch.tensor(transitions.state, device=self.device, dtype=torch.float32) nonFinalMask = torch.tensor(tuple( map(lambda s: s is not None, transitions.next_state)), device=self.device, dtype=torch.uint8) nonFinalNextState = torch.tensor( [s for s in transitions.next_state if s is not None], device=self.device, dtype=torch.float32) return state, nonFinalMask, nonFinalNextState, action, reward def update_net(self, state, action, nextState, reward, info): # first store memory self.store_experience(state, action, nextState, reward, info) # prepare mini-batch if len(self.memory) < self.trainBatchSize: return transitions_raw = self.memory.sample(self.trainBatchSize) self.update_net_on_transitions(transitions_raw) self.copy_nets() self.learnStepCounter += 1 def copy_nets(self): # update networks if self.learnStepCounter % self.policyUpdateFreq == 0: # update target networks for target_param, param in zip(self.actorNet_target.parameters(), self.actorNet.parameters()): target_param.data.copy_(param.data * self.tau + target_param.data * (1.0 - self.tau)) for target_param, param in zip(self.criticNet_target.parameters(), self.criticNet.parameters()): target_param.data.copy_(param.data * self.tau + target_param.data * (1.0 - self.tau)) def update_net_on_transitions(self, transitions_raw): state, nonFinalMask, nonFinalNextState, action, reward = self.prepare_minibatch( transitions_raw) # Critic loss QValues = self.criticNet.forward(state, action).squeeze() QNext = torch.zeros(self.trainBatchSize, device=self.device, dtype=torch.float32) if len(nonFinalNextState): # next action is calculated using target actor network next_actions = self.actorNet_target.forward(nonFinalNextState) QNext[nonFinalMask] = self.criticNet_target.forward( nonFinalNextState, next_actions.detach()).squeeze() targetValues = reward + self.gamma * QNext critic_loss = self.netLossFunc(QValues, targetValues) self.critic_optimizer.zero_grad() critic_loss.backward() if self.netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.criticNet.parameters(), self.netGradClip) self.critic_optimizer.step() # Actor loss # we try to maximize criticNet output(which is state value) # update networks if self.learnStepCounter % self.policyUpdateFreq == 0: policy_loss = -self.criticNet.forward( state, self.actorNet.forward(state)).mean() self.actor_optimizer.zero_grad() policy_loss.backward() if self.netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.actorNet.parameters(), self.netGradClip) self.actor_optimizer.step() if self.globalStepCount % self.lossRecordStep == 0: self.losses.append([ self.globalStepCount, self.epIdx, critic_loss.item(), policy_loss.item() ]) def work_before_step(self, state): pass def train(self): runningAvgEpisodeReward = 0.0 if len(self.rewards) > 0: runningAvgEpisodeReward = self.rewards[-1][-1] for trainStepCount in range(self.trainStep): print("episode index:" + str(self.epIdx)) state = self.env.reset() done = False rewardSum = 0 for stepCount in range(self.episodeLength): # any work to be done before select actions self.work_before_step(state) action = self.select_action(self.actorNet, state, noiseFlag=True) nextState, reward, done, info = self.env.step(action) if stepCount == 0: print("at step 0:") print(info) if done: nextState = None # learn the transition self.update_net(state, action, nextState, reward, info) state = nextState rewardSum += reward * pow(self.gamma, stepCount) self.globalStepCount += 1 if self.verbose: print('action: ' + str(action)) print('state:') print(nextState) print('reward:') print(reward) print('info') print(info) if done: break runningAvgEpisodeReward = (runningAvgEpisodeReward * self.epIdx + rewardSum) / (self.epIdx + 1) print("done in step count: {}".format(stepCount)) print("reward sum = " + str(rewardSum)) print("running average episode reward sum: {}".format( runningAvgEpisodeReward)) print(info) self.rewards.append([ self.epIdx, stepCount, self.globalStepCount, rewardSum, runningAvgEpisodeReward ]) if self.config['logFlag'] and self.epIdx % self.config[ 'logFrequency'] == 0: self.save_checkpoint() self.epIdx += 1 self.save_all() def saveLosses(self, fileName): np.savetxt(fileName, np.array(self.losses), fmt='%.5f', delimiter='\t') def saveRewards(self, fileName): np.savetxt(fileName, np.array(self.rewards), fmt='%.5f', delimiter='\t') def loadLosses(self, fileName): self.losses = np.genfromtxt(fileName).tolist() def loadRewards(self, fileName): self.rewards = np.genfromtxt(fileName).tolist() def save_all(self): prefix = self.dirName + self.identifier + 'Finalepoch' + str( self.epIdx) self.saveLosses(prefix + '_loss.txt') self.saveRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'wb') as file: pickle.dump(self.memory, file) torch.save( { 'epoch': self.epIdx, 'globalStep': self.globalStepCount, 'actorNet_state_dict': self.actorNet.state_dict(), 'criticNet_state_dict': self.criticNet.state_dict(), 'actor_optimizer_state_dict': self.actor_optimizer.state_dict(), 'critic_optimizer_state_dict': self.critic_optimizer.state_dict() }, prefix + '_checkpoint.pt') def save_checkpoint(self): prefix = self.dirName + self.identifier + 'Epoch' + str(self.epIdx) self.saveLosses(prefix + '_loss.txt') self.saveRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'wb') as file: pickle.dump(self.memory, file) torch.save( { 'epoch': self.epIdx, 'globalStep': self.globalStepCount, 'actorNet_state_dict': self.actorNet.state_dict(), 'criticNet_state_dict': self.criticNet.state_dict(), 'actor_optimizer_state_dict': self.actor_optimizer.state_dict(), 'critic_optimizer_state_dict': self.critic_optimizer.state_dict() }, prefix + '_checkpoint.pt') def load_checkpoint(self, prefix): self.loadLosses(prefix + '_loss.txt') self.loadRewards(prefix + '_reward.txt') with open(prefix + '_memory.pickle', 'rb') as file: self.memory = pickle.load(file) checkpoint = torch.load(prefix + '_checkpoint.pt') self.epIdx = checkpoint['epoch'] self.globalStepCount = checkpoint['globalStep'] self.actorNet.load_state_dict(checkpoint['actorNet_state_dict']) self.actorNet_target.load_state_dict(checkpoint['actorNet_state_dict']) self.criticNet.load_state_dict(checkpoint['criticNet_state_dict']) self.criticNet_target.load_state_dict( checkpoint['criticNet_state_dict']) self.actor_optimizer.load_state_dict( checkpoint['actor_optimizer_state_dict']) self.critic_optimizer.load_state_dict( checkpoint['critic_optimizer_state_dict'])
def init_memory(self): if self.memoryOption != 'natural': raise NotImplementedError self.memory = ReplayMemory(self.memoryCapacity)
class NAFAgent(DQNAgent): def __init__(self, config, policyNet, targetNet, env, optimizer, netLossFunc, nbAction, stateProcessor = None, experienceProcessor=None): super(NAFAgent, self).__init__(config, policyNet, targetNet, env, optimizer, netLossFunc, nbAction, stateProcessor, experienceProcessor) self.init_memory() def init_memory(self): if self.memoryOption != 'natural': raise NotImplementedError self.memory = ReplayMemory(self.memoryCapacity) def read_config(self): super(NAFAgent, self).read_config() # read additional parameters self.tau = self.config['tau'] def work_At_Episode_Begin(self): pass def work_before_step(self, state=None): self.epsThreshold = self.epsilon_by_step(self.globalStepCount) def train_one_episode(self): print("episode index:" + str(self.epIdx)) state = self.env.reset() done = False rewardSum = 0 # any work to be done at episode begin self.work_At_Episode_Begin() for stepCount in range(self.episodeLength): # any work to be done before select actions self.work_before_step(state) action = self.select_action(self.policyNet, state, noiseFlag=True) nextState, reward, done, info = self.env.step(action) if stepCount == 0: print("at step 0:") print(info) if done: nextState = None # learn the transition self.update_net(state, action, nextState, reward, info) state = nextState rewardSum += reward * pow(self.gamma, stepCount) self.globalStepCount += 1 if self.verbose: print('action: ' + str(action)) print('state:') print(nextState) print('reward:') print(reward) print('info') print(info) if done: break self.runningAvgEpisodeReward = (self.runningAvgEpisodeReward * self.epIdx + rewardSum) / (self.epIdx + 1) print("done in step count: {}".format(stepCount)) print("reward sum = " + str(rewardSum)) print("running average episode reward sum: {}".format(self.runningAvgEpisodeReward)) print(info) self.rewards.append([self.epIdx, stepCount, self.globalStepCount, rewardSum, self.runningAvgEpisodeReward]) if self.config['logFlag'] and self.epIdx % self.config['logFrequency'] == 0: self.save_checkpoint() self.epIdx += 1 return stepCount, rewardSum def train(self): if len(self.rewards) > 0: self.runningAvgEpisodeReward = self.rewards[-1][-1] for trainStepCount in range(self.trainStep): self.train_one_episode() self.save_all() def store_experience(self, state, action, nextState, reward, info): if self.experienceProcessor is not None: state, action, nextState, reward = self.experienceProcessor(state, action, nextState, reward, info) transition = Transition(state, action, nextState, reward) self.memory.push(transition) def update_net(self, state, action, nextState, reward, info): # first store memory self.store_experience(state, action, nextState, reward, info) if self.hindSightER and nextState is not None and self.globalStepCount % self.hindSightERFreq == 0: stateNew, actionNew, nextStateNew, rewardNew = self.env.getHindSightExperience(state, action, nextState, info) if stateNew is not None: self.store_experience(stateNew, actionNew, nextStateNew, rewardNew, info) if self.priorityMemoryOption: if len(self.memory) < self.config['memoryCapacity']: return else: if len(self.memory) < self.trainBatchSize: return # update net with specified frequency if self.globalStepCount % self.netUpdateFrequency == 0: # sample experience for nStep in range(self.netUpdateStep): info = {} if self.priorityMemoryOption: transitions_raw, b_idx, ISWeights = self.memory.sample(self.trainBatchSize) info['batchIdx'] = b_idx info['ISWeights'] = torch.from_numpy(ISWeights.astype(np.float32)).to(self.device) else: transitions_raw= self.memory.sample(self.trainBatchSize) loss = self.update_net_on_transitions(transitions_raw, self.netLossFunc, 1, updateOption=self.netUpdateOption, netGradClip=self.netGradClip, info=info) if self.globalStepCount % self.lossRecordStep == 0: self.losses.append([self.globalStepCount, self.epIdx, loss]) if self.learnStepCounter % self.targetNetUpdateStep == 0: self.targetNet.load_state_dict(self.policyNet.state_dict()) self.learnStepCounter += 1 def prepare_minibatch(self, transitions_raw): # first store memory transitions = Transition(*zip(*transitions_raw)) action = torch.tensor(transitions.action, device=self.device, dtype=torch.float32) # shape(batch, numActions) reward = torch.tensor(transitions.reward, device=self.device, dtype=torch.float32) # shape(batch) # for some env, the output state requires further processing before feeding to neural network if self.stateProcessor is not None: state, _ = self.stateProcessor(transitions.state, self.device) nonFinalNextState, nonFinalMask = self.stateProcessor(transitions.next_state, self.device) else: state = torch.tensor(transitions.state, device=self.device, dtype=torch.float32) nonFinalMask = torch.tensor(tuple(map(lambda s: s is not None, transitions.next_state)), device=self.device, dtype=torch.uint8) nonFinalNextState = torch.tensor([s for s in transitions.next_state if s is not None], device=self.device, dtype=torch.float32) return state, nonFinalMask, nonFinalNextState, action, reward def select_action(self, net=None, state=None, noiseFlag = True): if net is None: net = self.policyNet with torch.no_grad(): # self.policyNet(torch.from_numpy(state.astype(np.float32)).unsqueeze(0)) # here state[np.newaxis,:] is to add a batch dimension if self.stateProcessor is not None: state, _ = self.stateProcessor([state], self.device) action = net.select_action(state, noiseFlag) else: stateTorch = torch.from_numpy(np.array(state[np.newaxis, :], dtype=np.float32)) action = net.select_action(stateTorch.to(self.device), noiseFlag) return action.cpu().data.numpy()[0] def update_net_on_transitions(self, transitions_raw, loss_fun, gradientStep = 1, updateOption='policyNet', netGradClip=None, info=None): # order the data # convert transition list to torch tensors # use trick from https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html # https://stackoverflow.com/questions/19339/transpose-unzip-function-inverse-of-zip/19343#19343 state, nonFinalMask, nonFinalNextState, action, reward = self.prepare_minibatch(transitions_raw) for step in range(gradientStep): # calculate Qvalues based on selected action batch QValues = self.policyNet.eval_Q_value(state, action).squeeze() # Here we detach because we do not want gradient flow from target values to net parameters QNext = torch.zeros(self.trainBatchSize, device=self.device, dtype=torch.float32) QNext[nonFinalMask] = self.targetNet.eval_state_value(nonFinalNextState).squeeze().detach() targetValues = reward + self.gamma * QNext # Compute loss loss_single = loss_fun(QValues, targetValues) loss = torch.mean(loss_single) # Optimize the model # zero gradient self.optimizer.zero_grad() loss.backward() if netGradClip is not None: torch.nn.utils.clip_grad_norm_(self.policyNet.parameters(), netGradClip) self.optimizer.step() return loss.item()