• 用深度强化学习来玩Chrome小恐龙快跑

    目录

    实机演示

    代码实现

    实机演示

    用深度强化学习来玩Chrome小恐龙快跑

    代码实现

    1. import os
    2. import cv2
    3. from pygame import RLEACCEL
    4. from pygame.image import load
    5. from pygame.sprite import Sprite, Group, collide_mask
    6. from pygame import Rect, init, time, display, mixer, transform, Surface
    7. from pygame.surfarray import array3d
    8. import torch
    9. from random import randrange, choice
    10. import numpy as np
    11.  
    12. mixer.pre_init(44100, -16, 2, 2048)
    13. init()
    14.  
    15. scr_size = (width, height) = (600, 150)
    16. FPS = 60
    17. gravity = 0.6
    18.  
    19. black = (0, 0, 0)
    20. white = (255, 255, 255)
    21. background_col = (235, 235, 235)
    22.  
    23. high_score = 0
    24.  
    25. screen = display.set_mode(scr_size)
    26. clock = time.Clock()
    27. display.set_caption("T-Rex Rush")
    28.  
    29.  
    30. def load_image(
    31.         name,
    32.         sizex=-1,
    33.         sizey=-1,
    34.         colorkey=None,
    35. ):
    36.     fullname = os.path.join("assets/sprites", name)
    37.     image = load(fullname)
    38.     image = image.convert()
    39.     if colorkey is not None:
    40.         if colorkey is -1:
    41.             colorkey = image.get_at((0, 0))
    42.         image.set_colorkey(colorkey, RLEACCEL)
    43.  
    44.     if sizex != -1 or sizey != -1:
    45.         image = transform.scale(image, (sizex, sizey))
    46.  
    47.     return (image, image.get_rect())
    48.  
    49.  
    50. def load_sprite_sheet(
    51.         sheetname,
    52.         nx,
    53.         ny,
    54.         scalex=-1,
    55.         scaley=-1,
    56.         colorkey=None,
    57. ):
    58.     fullname = os.path.join("assets/sprites", sheetname)
    59.     sheet = load(fullname)
    60.     sheet = sheet.convert()
    61.  
    62.     sheet_rect = sheet.get_rect()
    63.  
    64.     sprites = []
    65.  
    66.     sizey = sheet_rect.height / ny
    67.     if isinstance(nx, int):
    68.         sizex = sheet_rect.width / nx
    69.         for i in range(0, ny):
    70.             for j in range(0, nx):
    71.                 rect = Rect((j * sizex, i * sizey, sizex, sizey))
    72.                 image = Surface(rect.size)
    73.                 image = image.convert()
    74.                 image.blit(sheet, (0, 0), rect)
    75.  
    76.                 if colorkey is not None:
    77.                     if colorkey is -1:
    78.                         colorkey = image.get_at((0, 0))
    79.                     image.set_colorkey(colorkey, RLEACCEL)
    80.  
    81.                 if scalex != -1 or scaley != -1:
    82.                     image = transform.scale(image, (scalex, scaley))
    83.  
    84.                 sprites.append(image)
    85.  
    86.     else:  #list
    87.         sizex_ls = [sheet_rect.width / i_nx for i_nx in nx]
    88.         for i in range(0, ny):
    89.             for i_nx, sizex, i_scalex in zip(nx, sizex_ls, scalex):
    90.                 for j in range(0, i_nx):
    91.                     rect = Rect((j * sizex, i * sizey, sizex, sizey))
    92.                     image = Surface(rect.size)
    93.                     image = image.convert()
    94.                     image.blit(sheet, (0, 0), rect)
    95.  
    96.                     if colorkey is not None:
    97.                         if colorkey is -1:
    98.                             colorkey = image.get_at((0, 0))
    99.                         image.set_colorkey(colorkey, RLEACCEL)
    100.  
    101.                     if i_scalex != -1 or scaley != -1:
    102.                         image = transform.scale(image, (i_scalex, scaley))
    103.  
    104.                     sprites.append(image)
    105.  
    106.     sprite_rect = sprites[0].get_rect()
    107.  
    108.     return sprites, sprite_rect
    109.  
    110.  
    111. def extractDigits(number):
    112.     if number > -1:
    113.         digits = []
    114.         i = 0
    115.         while (number / 10 != 0):
    116.             digits.append(number % 10)
    117.             number = int(number / 10)
    118.  
    119.         digits.append(number % 10)
    120.         for i in range(len(digits), 5):
    121.             digits.append(0)
    122.         digits.reverse()
    123.         return digits
    124.  
    125.  
    126. def pre_processing(image, w=84, h=84):
    127.     image = image[:300, :, :]
    128.     # cv2.imwrite("ori.jpg", image)
    129.     image = cv2.cvtColor(cv2.resize(image, (w, h)), cv2.COLOR_BGR2GRAY)
    130.     # cv2.imwrite("color.jpg", image)
    131.     _, image = cv2.threshold(image, 127, 255, cv2.THRESH_BINARY)
    132.     # cv2.imwrite("bw.jpg", image)
    133.  
    134.     return image[None, :, :].astype(np.float32)
    135.  
    136.  
    137. class Dino():
    138.     def __init__(self, sizex=-1, sizey=-1):
    139.         self.images, self.rect = load_sprite_sheet("dino.png", 5, 1, sizex, sizey, -1)
    140.         self.images1, self.rect1 = load_sprite_sheet("dino_ducking.png", 2, 1, 59, sizey, -1)
    141.         self.rect.bottom = int(0.98 * height)
    142.         self.rect.left = width / 15
    143.         self.image = self.images[0]
    144.         self.index = 0
    145.         self.counter = 0
    146.         self.score = 0
    147.         self.isJumping = False
    148.         self.isDead = False
    149.         self.isDucking = False
    150.         self.isBlinking = False
    151.         self.movement = [0, 0]
    152.         self.jumpSpeed = 11.5
    153.  
    154.         self.stand_pos_width = self.rect.width
    155.         self.duck_pos_width = self.rect1.width
    156.  
    157.     def draw(self):
    158.         screen.blit(self.image, self.rect)
    159.  
    160.     def checkbounds(self):
    161.         if self.rect.bottom > int(0.98 * height):
    162.             self.rect.bottom = int(0.98 * height)
    163.             self.isJumping = False
    164.  
    165.     def update(self):
    166.         if self.isJumping:
    167.             self.movement[1] = self.movement[1] + gravity
    168.  
    169.         if self.isJumping:
    170.             self.index = 0
    171.         elif self.isBlinking:
    172.             if self.index == 0:
    173.                 if self.counter % 400 == 399:
    174.                     self.index = (self.index + 1) % 2
    175.             else:
    176.                 if self.counter % 20 == 19:
    177.                     self.index = (self.index + 1) % 2
    178.  
    179.         elif self.isDucking:
    180.             if self.counter % 5 == 0:
    181.                 self.index = (self.index + 1) % 2
    182.         else:
    183.             if self.counter % 5 == 0:
    184.                 self.index = (self.index + 1) % 2 + 2
    185.  
    186.         if self.isDead:
    187.             self.index = 4
    188.  
    189.         if not self.isDucking:
    190.             self.image = self.images[self.index]
    191.             self.rect.width = self.stand_pos_width
    192.         else:
    193.             self.image = self.images1[(self.index) % 2]
    194.             self.rect.width = self.duck_pos_width
    195.  
    196.         self.rect = self.rect.move(self.movement)
    197.         self.checkbounds()
    198.  
    199.         if not self.isDead and self.counter % 7 == 6 and self.isBlinking == False:
    200.             self.score += 1
    201.  
    202.         self.counter = (self.counter + 1)
    203.  
    204.  
    205. class Cactus(Sprite):
    206.     def __init__(self, speed=5, sizex=-1, sizey=-1):
    207.         Sprite.__init__(self, self.containers)
    208.         self.images, self.rect = load_sprite_sheet("cacti-small.png", [2, 3, 6], 1, sizex, sizey, -1)
    209.         self.rect.bottom = int(0.98 * height)
    210.         self.rect.left = width + self.rect.width
    211.         self.image = self.images[randrange(0, 11)]
    212.         self.movement = [-1 * speed, 0]
    213.  
    214.     def draw(self):
    215.         screen.blit(self.image, self.rect)
    216.  
    217.     def update(self):
    218.         self.rect = self.rect.move(self.movement)
    219.  
    220.         if self.rect.right < 0:
    221.             self.kill()
    222.  
    223.  
    224. class Ptera(Sprite):
    225.     def __init__(self, speed=5, sizex=-1, sizey=-1):
    226.         Sprite.__init__(self, self.containers)
    227.         self.images, self.rect = load_sprite_sheet("ptera.png", 2, 1, sizex, sizey, -1)
    228.         self.ptera_height = [height * 0.82, height * 0.75, height * 0.60, height * 0.48]
    229.         self.rect.centery = self.ptera_height[randrange(0, 4)]
    230.         self.rect.left = width + self.rect.width
    231.         self.image = self.images[0]
    232.         self.movement = [-1 * speed, 0]
    233.         self.index = 0
    234.         self.counter = 0
    235.  
    236.     def draw(self):
    237.         screen.blit(self.image, self.rect)
    238.  
    239.     def update(self):
    240.         if self.counter % 10 == 0:
    241.             self.index = (self.index + 1) % 2
    242.         self.image = self.images[self.index]
    243.         self.rect = self.rect.move(self.movement)
    244.         self.counter = (self.counter + 1)
    245.         if self.rect.right < 0:
    246.             self.kill()
    247.  
    248.  
    249. class Ground():
    250.     def __init__(self, speed=-5):
    251.         self.image, self.rect = load_image("ground.png", -1, -1, -1)
    252.         self.image1, self.rect1 = load_image("ground.png", -1, -1, -1)
    253.         self.rect.bottom = height
    254.         self.rect1.bottom = height
    255.         self.rect1.left = self.rect.right
    256.         self.speed = speed
    257.  
    258.     def draw(self):
    259.         screen.blit(self.image, self.rect)
    260.         screen.blit(self.image1, self.rect1)
    261.  
    262.     def update(self):
    263.         self.rect.left += self.speed
    264.         self.rect1.left += self.speed
    265.  
    266.         if self.rect.right < 0:
    267.             self.rect.left = self.rect1.right
    268.  
    269.         if self.rect1.right < 0:
    270.             self.rect1.left = self.rect.right
    271.  
    272.  
    273. class Cloud(Sprite):
    274.     def __init__(self, x, y):
    275.         Sprite.__init__(self, self.containers)
    276.         self.image, self.rect = load_image("cloud.png", int(90 * 30 / 42), 30, -1)
    277.         self.speed = 1
    278.         self.rect.left = x
    279.         self.rect.top = y
    280.         self.movement = [-1 * self.speed, 0]
    281.  
    282.     def draw(self):
    283.         screen.blit(self.image, self.rect)
    284.  
    285.     def update(self):
    286.         self.rect = self.rect.move(self.movement)
    287.         if self.rect.right < 0:
    288.             self.kill()
    289.  
    290.  
    291. class Scoreboard():
    292.     def __init__(self, x=-1, y=-1):
    293.         self.score = 0
    294.         self.tempimages, self.temprect = load_sprite_sheet("numbers.png", 12, 1, 11, int(11 * 6 / 5), -1)
    295.         self.image = Surface((55, int(11 * 6 / 5)))
    296.         self.rect = self.image.get_rect()
    297.         if x == -1:
    298.             self.rect.left = width * 0.89
    299.         else:
    300.             self.rect.left = x
    301.         if y == -1:
    302.             self.rect.top = height * 0.1
    303.         else:
    304.             self.rect.top = y
    305.  
    306.     def draw(self):
    307.         screen.blit(self.image, self.rect)
    308.  
    309.     def update(self, score):
    310.         score_digits = extractDigits(score)
    311.         self.image.fill(background_col)
    312.         if len(score_digits) == 6:
    313.             score_digits = score_digits[1:]
    314.         for s in score_digits:
    315.             self.image.blit(self.tempimages[s], self.temprect)
    316.             self.temprect.left += self.temprect.width
    317.         self.temprect.left = 0
    318.  
    319.  
    320. class ChromeDino(object):
    321.     def __init__(self):
    322.         self.gamespeed = 5
    323.         self.gameOver = False
    324.         self.gameQuit = False
    325.         self.playerDino = Dino(44, 47)
    326.         self.new_ground = Ground(-1 * self.gamespeed)
    327.         self.scb = Scoreboard()
    328.         self.highsc = Scoreboard(width * 0.78)
    329.         self.counter = 0
    330.  
    331.         self.cacti = Group()
    332.         self.pteras = Group()
    333.         self.clouds = Group()
    334.         self.last_obstacle = Group()
    335.  
    336.         Cactus.containers = self.cacti
    337.         Ptera.containers = self.pteras
    338.         Cloud.containers = self.clouds
    339.  
    340.         self.retbutton_image, self.retbutton_rect = load_image("replay_button.png", 35, 31, -1)
    341.         self.gameover_image, self.gameover_rect = load_image("game_over.png", 190, 11, -1)
    342.  
    343.         self.temp_images, self.temp_rect = load_sprite_sheet("numbers.png", 12, 1, 11, int(11 * 6 / 5), -1)
    344.         self.HI_image = Surface((22, int(11 * 6 / 5)))
    345.         self.HI_rect = self.HI_image.get_rect()
    346.         self.HI_image.fill(background_col)
    347.         self.HI_image.blit(self.temp_images[10], self.temp_rect)
    348.         self.temp_rect.left += self.temp_rect.width
    349.         self.HI_image.blit(self.temp_images[11], self.temp_rect)
    350.         self.HI_rect.top = height * 0.1
    351.         self.HI_rect.left = width * 0.73
    352.  
    353.     def step(self, action, record=False):  # 0: Do nothing. 1: Jump. 2: Duck
    354.         reward = 0.1
    355.         if action == 0:
    356.             reward += 0.01
    357.             self.playerDino.isDucking = False
    358.         elif action == 1:
    359.             self.playerDino.isDucking = False
    360.             if self.playerDino.rect.bottom == int(0.98 * height):
    361.                 self.playerDino.isJumping = True
    362.                 self.playerDino.movement[1] = -1 * self.playerDino.jumpSpeed
    363.  
    364.         elif action == 2:
    365.             if not (self.playerDino.isJumping and self.playerDino.isDead) and self.playerDino.rect.bottom == int(
    366.                     0.98 * height):
    367.                 self.playerDino.isDucking = True
    368.  
    369.         for c in self.cacti:
    370.             c.movement[0] = -1 * self.gamespeed
    371.             if collide_mask(self.playerDino, c):
    372.                 self.playerDino.isDead = True
    373.                 reward = -1
    374.                 break
    375.             else:
    376.                 if c.rect.right < self.playerDino.rect.left < c.rect.right + self.gamespeed + 1:
    377.                     reward = 1
    378.                     break
    379.  
    380.         for p in self.pteras:
    381.             p.movement[0] = -1 * self.gamespeed
    382.             if collide_mask(self.playerDino, p):
    383.                 self.playerDino.isDead = True
    384.                 reward = -1
    385.                 break
    386.             else:
    387.                 if p.rect.right < self.playerDino.rect.left < p.rect.right + self.gamespeed + 1:
    388.                     reward = 1
    389.                     break
    390.  
    391.         if len(self.cacti) < 2:
    392.             if len(self.cacti) == 0 and len(self.pteras) == 0:
    393.                 self.last_obstacle.empty()
    394.                 self.last_obstacle.add(Cactus(self.gamespeed, [60, 40, 20], choice([40, 45, 50])))
    395.             else:
    396.                 for l in self.last_obstacle:
    397.                     if l.rect.right < width * 0.7 and randrange(0, 50) == 10:
    398.                         self.last_obstacle.empty()
    399.                         self.last_obstacle.add(Cactus(self.gamespeed, [60, 40, 20], choice([40, 45, 50])))
    400.  
    401.         # if len(self.pteras) == 0 and randrange(0, 200) == 10 and self.counter > 500:
    402.         if len(self.pteras) == 0 and len(self.cacti) < 2 and randrange(0, 50) == 10 and self.counter > 500:
    403.             for l in self.last_obstacle:
    404.                 if l.rect.right < width * 0.8:
    405.                     self.last_obstacle.empty()
    406.                     self.last_obstacle.add(Ptera(self.gamespeed, 46, 40))
    407.  
    408.         if len(self.clouds) < 5 and randrange(0, 300) == 10:
    409.             Cloud(width, randrange(height / 5, height / 2))
    410.  
    411.         self.playerDino.update()
    412.         self.cacti.update()
    413.         self.pteras.update()
    414.         self.clouds.update()
    415.         self.new_ground.update()
    416.         self.scb.update(self.playerDino.score)
    417.  
    418.         state = display.get_surface()
    419.         screen.fill(background_col)
    420.         self.new_ground.draw()
    421.         self.clouds.draw(screen)
    422.         self.scb.draw()
    423.         self.cacti.draw(screen)
    424.         self.pteras.draw(screen)
    425.         self.playerDino.draw()
    426.  
    427.         display.update()
    428.         clock.tick(FPS)
    429.  
    430.         if self.playerDino.isDead:
    431.             self.gameOver = True
    432.  
    433.         self.counter = (self.counter + 1)
    434.  
    435.         if self.gameOver:
    436.             self.__init__()
    437.  
    438.         state = array3d(state)
    439.         if record:
    440.             return torch.from_numpy(pre_processing(state)), np.transpose(
    441.                 cv2.cvtColor(state, cv2.COLOR_RGB2BGR), (1, 0, 2)), reward, not (reward > 0)
    442.         else:
    443.             return torch.from_numpy(pre_processing(state)), reward, not (reward > 0)
    1. import torch.nn as nn
    2.  
    3. class DeepQNetwork(nn.Module):
    4.     def __init__(self):
    5.         super(DeepQNetwork, self).__init__()
    6.  
    7.         self.conv1 = nn.Sequential(nn.Conv2d(4, 32, kernel_size=8, stride=4), nn.ReLU(inplace=True))
    8.         self.conv2 = nn.Sequential(nn.Conv2d(32, 64, kernel_size=4, stride=2), nn.ReLU(inplace=True))
    9.         self.conv3 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, stride=1), nn.ReLU(inplace=True))
    10.  
    11.         self.fc1 = nn.Sequential(nn.Linear(7 * 7 * 64, 512), nn.ReLU(inplace=True))
    12.         self.fc2 = nn.Linear(512, 3)
    13.         self._initialize_weights()
    14.  
    15.     def _initialize_weights(self):
    16.         for m in self.modules():
    17.             if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
    18.                 nn.init.uniform_(m.weight, -0.01, 0.01)
    19.                 nn.init.constant_(m.bias, 0)
    20.  
    21.     def forward(self, input):
    22.         output = self.conv1(input)
    23.         output = self.conv2(output)
    24.         output = self.conv3(output)
    25.         output = output.view(output.size(0), -1)
    26.         output = self.fc1(output)
    27.         output = self.fc2(output)
    28.  
    29.         return output
    1. import argparse
    2. import torch
    3.  
    4. from src.model import DeepQNetwork
    5. from src.env import ChromeDino
    6. import cv2
    7.  
    8.  
    9. def get_args():
    10.     parser = argparse.ArgumentParser(
    11.         """Implementation of Deep Q Network to play Chrome Dino""")
    12.     parser.add_argument("--saved_path", type=str, default="trained_models")
    13.     parser.add_argument("--fps", type=int, default=60, help="frames per second")
    14.     parser.add_argument("--output", type=str, default="output/chrome_dino.mp4", help="the path to output video")
    15.  
    16.     args = parser.parse_args()
    17.     return args
    18.  
    19.  
    20. def q_test(opt):
    21.     if torch.cuda.is_available():
    22.         torch.cuda.manual_seed(123)
    23.     else:
    24.         torch.manual_seed(123)
    25.     model = DeepQNetwork()
    26.     checkpoint_path = "{}/chrome_dino.pth".format(opt.saved_path)
    27.     checkpoint = torch.load(checkpoint_path)
    28.     model.load_state_dict(checkpoint["model_state_dict"])
    29.     model.eval()
    30.     env = ChromeDino()
    31.     state, raw_state, _, _ = env.step(0, True)
    32.     state = torch.cat(tuple(state for _ in range(4)))[None, :, :, :]
    33.     if torch.cuda.is_available():
    34.         model.cuda()
    35.         state = state.cuda()
    36.     out = cv2.VideoWriter(opt.output, cv2.VideoWriter_fourcc(*"MJPG"), opt.fps, (600, 150))
    37.     done = False
    38.     while not done:
    39.         prediction = model(state)[0]
    40.         action = torch.argmax(prediction).item()
    41.         next_state, raw_next_state, reward, done = env.step(action, True)
    42.         out.write(raw_next_state)
    43.         if torch.cuda.is_available():
    44.             next_state = next_state.cuda()
    45.         next_state = torch.cat((state[0, 1:, :, :], next_state))[None, :, :, :]
    46.         state = next_state
    47.  
    48.  
    49.  
    50. if __name__ == "__main__":
    51.     opt = get_args()
    52.     q_test(opt)
    1. import argparse
    2. import os
    3. from random import random, randint, sample
    4. import pickle
    5. import numpy as np
    6. import torch
    7. import torch.nn as nn
    8.  
    9. from src.model import DeepQNetwork
    10. from src.env import ChromeDino
    11.  
    12.  
    13. def get_args():
    14.     parser = argparse.ArgumentParser(
    15.         """Implementation of Deep Q Network to play Chrome Dino""")
    16.     parser.add_argument("--batch_size", type=int, default=64, help="The number of images per batch")
    17.     parser.add_argument("--optimizer", type=str, choices=["sgd", "adam"], default="adam")
    18.     parser.add_argument("--lr", type=float, default=1e-4)
    19.     parser.add_argument("--gamma", type=float, default=0.99)
    20.     parser.add_argument("--initial_epsilon", type=float, default=0.1)
    21.     parser.add_argument("--final_epsilon", type=float, default=1e-4)
    22.     parser.add_argument("--num_decay_iters", type=float, default=2000000)
    23.     parser.add_argument("--num_iters", type=int, default=2000000)
    24.     parser.add_argument("--replay_memory_size", type=int, default=50000,
    25.                         help="Number of epoches between testing phases")
    26.     parser.add_argument("--saved_folder", type=str, default="trained_models")
    27.  
    28.     args = parser.parse_args()
    29.     return args
    30.  
    31.  
    32. def train(opt):
    33.     if torch.cuda.is_available():
    34.         torch.cuda.manual_seed(123)
    35.     else:
    36.         torch.manual_seed(123)
    37.     model = DeepQNetwork()
    38.     if torch.cuda.is_available():
    39.         model.cuda()
    40.     optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr)
    41.     if not os.path.isdir(opt.saved_folder):
    42.         os.makedirs(opt.saved_folder)
    43.     checkpoint_path = os.path.join(opt.saved_folder, "chrome_dino.pth")
    44.     memory_path = os.path.join(opt.saved_folder, "replay_memory.pkl")
    45.     if os.path.isfile(checkpoint_path):
    46.         checkpoint = torch.load(checkpoint_path)
    47.         iter = checkpoint["iter"] + 1
    48.         model.load_state_dict(checkpoint["model_state_dict"])
    49.         optimizer.load_state_dict(checkpoint["optimizer"])
    50.         print("Load trained model from iteration {}".format(iter))
    51.     else:
    52.         iter = 0
    53.     if os.path.isfile(memory_path):
    54.         with open(memory_path, "rb") as f:
    55.             replay_memory = pickle.load(f)
    56.         print("Load replay memory")
    57.     else:
    58.         replay_memory = []
    59.     criterion = nn.MSELoss()
    60.     env = ChromeDino()
    61.     state, _, _ = env.step(0)
    62.     state = torch.cat(tuple(state for _ in range(4)))[None, :, :, :]
    63.     while iter < opt.num_iters:
    64.         if torch.cuda.is_available():
    65.             prediction = model(state.cuda())[0]
    66.         else:
    67.             prediction = model(state)[0]
    68.         # Exploration or exploitation
    69.         epsilon = opt.final_epsilon + (
    70.                 max(opt.num_decay_iters - iter, 0) * (opt.initial_epsilon - opt.final_epsilon) / opt.num_decay_iters)
    71.         u = random()
    72.         random_action = u <= epsilon
    73.         if random_action:
    74.             action = randint(0, 2)
    75.         else:
    76.             action = torch.argmax(prediction).item()
    77.  
    78.         next_state, reward, done = env.step(action)
    79.         next_state = torch.cat((state[0, 1:, :, :], next_state))[None, :, :, :]
    80.         replay_memory.append([state, action, reward, next_state, done])
    81.         if len(replay_memory) > opt.replay_memory_size:
    82.             del replay_memory[0]
    83.         batch = sample(replay_memory, min(len(replay_memory), opt.batch_size))
    84.         state_batch, action_batch, reward_batch, next_state_batch, done_batch = zip(*batch)
    85.  
    86.         state_batch = torch.cat(tuple(state for state in state_batch))
    87.         action_batch = torch.from_numpy(
    88.             np.array([[1, 0, 0] if action == 0 else [0, 1, 0] if action == 1 else [0, 0, 1] for action in
    89.                       action_batch], dtype=np.float32))
    90.         reward_batch = torch.from_numpy(np.array(reward_batch, dtype=np.float32)[:, None])
    91.         next_state_batch = torch.cat(tuple(state for state in next_state_batch))
    92.  
    93.         if torch.cuda.is_available():
    94.             state_batch = state_batch.cuda()
    95.             action_batch = action_batch.cuda()
    96.             reward_batch = reward_batch.cuda()
    97.             next_state_batch = next_state_batch.cuda()
    98.         current_prediction_batch = model(state_batch)
    99.         next_prediction_batch = model(next_state_batch)
    100.  
    101.         y_batch = torch.cat(
    102.             tuple(reward if done else reward + opt.gamma * torch.max(prediction) for reward, done, prediction in
    103.                   zip(reward_batch, done_batch, next_prediction_batch)))
    104.  
    105.         q_value = torch.sum(current_prediction_batch * action_batch, dim=1)
    106.         optimizer.zero_grad()
    107.         loss = criterion(q_value, y_batch)
    108.         loss.backward()
    109.         optimizer.step()
    110.  
    111.         state = next_state
    112.         iter += 1
    113.         print("Iteration: {}/{}, Loss: {:.5f}, Epsilon {:.5f}, Reward: {}".format(
    114.             iter + 1,
    115.             opt.num_iters,
    116.             loss,
    117.             epsilon, reward))
    118.         if (iter + 1) % 50000 == 0:
    119.             checkpoint = {"iter": iter,
    120.                           "model_state_dict": model.state_dict(),
    121.                           "optimizer": optimizer.state_dict()}
    122.             torch.save(checkpoint, checkpoint_path)
    123.             with open(memory_path, "wb") as f:
    124.                 pickle.dump(replay_memory, f, protocol=pickle.HIGHEST_PROTOCOL)
    125.  
    126.  
    127. if __name__ == "__main__":
    128.     opt = get_args()
    129.     train(opt)

  • 相关阅读:
    这份SVN命令备忘清单，请查收
    深度学习入门（五十六）循环神经网络——循环神经网络RNN
    python+playwright 学习-81 page.expect_request()捕获网络请求
    arm服务器运行onlyoffice
    ffmpeg基础三：H264，从MP4文件获取(av_bsf_get_by_name(“h264_mp4toannexb“))和从TS流获取保存H264
    使用 React 和 MUI 创建多选 Checkbox 树组件
    【开源代码】-基于国民N32G43x系列MCU使用JLINK的开发组件工具-RTTViewer/logger/Client/调试打印
    什么是作业指导书sop？sop作业指导书是什么意思？
    探索人工智能领域——每日30个名词详解【day3】
    积加（跨境ERP）与金蝶云星空单据集成对接
  • 原文地址：https://blog.csdn.net/timberman666/article/details/132630661