【低照度图像增强系列（7）】RDDNet算法详解与代码实现（同济大学|ICME）

已于 2024-07-12 14:25:26 修改
阅读量1.9k 收藏 31
点赞数 29
分类专栏： 低照度图像增强 文章标签： 图像增强 计算机视觉 人工智能 深度学习
于 2024-05-24 15:00:10 首次发布
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/weixin_43334693/article/details/139169688
版权

前言

☀️ 在低照度场景下进行目标检测任务，常存在图像RGB特征信息少、提取特征困难、目标识别和定位精度低等问题，给检测带来一定的难度。

     🌻使用图像增强模块对原始图像进行画质提升，恢复各类图像信息，再使用目标检测网络对增强图像进行特定目标检测，有效提高检测的精确度。

      ⭐本专栏会介绍传统方法、Retinex、EnlightenGAN、SCI、Zero-DCE、IceNet、RRDNet、URetinex-Net等低照度图像增强算法。

👑完整代码已打包上传至资源→低照度图像增强代码汇总资源

目录

前言

🚀一、RDDNet介绍 

☀️1.1 RDDNet简介   

研究背景 

算法框架 

损失函数

🚀二、RDDNet核心代码

 ☀️2.1 网络模型—RRDNet.py

 ☀️2.2 损失函数—loss_functions.py

（1）重构损失——reconstruction_loss

（2）光照损失——illumination_smooth_loss

（3）反射损失——reflectance_smooth_loss

（4）噪声损失——noise_loss

  ☀️2.3 Retinex操作—pipline.py

🚀三、RDDNet代码复现

☀️3.1 环境配置

☀️3.2 运行过程

☀️3.3 运行效果

🚀一、RDDNet介绍 

学习资料：

☀️1.1 RDDNet简介   

RRDNet是同济大学在2020年提出来的一种新的三分支全卷积神经网络，认为图像由三部分构成：光照分量，反射分量，噪声分量。在没有pair对的情况下实现低光图像增强，通过对loss进行迭代来有效估计出噪声和恢复光照。 

研究背景 

  • 曝光不足的图像由于能见度差和黑暗中的潜在噪声，通常会出现严重的质量下降。
  • 现有的图像增强方法忽略了噪声，因此使用带噪声分量的Retinex模型作为基础。
  • 基于学习（数据驱动）的方法限制了模型的泛化能力，因此提出zero-shot的学习模式。

算法框架 

  1. 通过三分支网络把输入图像分解为反射图、光照图和噪声图三个分量。
  2. 通过Gamma变换调整光照图，再计算得到无噪声的反射图。
  3. 结合光照图和反射图，重构得到最终结果。 

损失函数

1. Retinex重构损失，取最大通道值作为初始光照图，用来约束光照图。在光照图的基础上约束反射图和噪声。

2. 纹理增强损失，通过平滑光照图可以帮助增强反射图的纹理。具体损失公式是带有权重的总变分损失，权重的设计规则是，梯度大的地方权重小，即权重与梯度成负相关即可，这里是将梯度经过高斯滤波放在分母。

3. 光照指导的噪声损失，根据噪声随着光照的变大而变大的假设，可以使用光照图来做权重指导，其次考虑两点：

（1）假定噪声范围限定

（2）通过平滑反射图来得到噪声，本身并没有直接得到噪声的损失，只是通过对反射图做总变分约束来去噪

🚀二、RDDNet核心代码

 代码框架如图所示：

（图片来源：【代码笔记】RRDNet 网络-CSDN博客 谢谢大佬！@chaiky） 

 ☀️2.1 网络模型—RRDNet.py

import torch
import torch.nn as nn

class RRDNet(nn.Module):
    def __init__(self):
        super(RRDNet, self).__init__()

 #----------- 1.illumination（光照估计）---------------------------#
        self.illumination_net = nn.Sequential(
            nn.Conv2d(3, 16, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(64, 32, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(32, 1, 3, 1, 1),

        )

 #----------- 2.reflectance（反射率估计）---------------------------#
        self.reflectance_net = nn.Sequential(
            nn.Conv2d(3, 16, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(64, 32, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(32, 3, 3, 1, 1)
        )

 #----------- 3.noise（噪声估计）---------------------------#
        self.noise_net = nn.Sequential(
            nn.Conv2d(3, 16, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(64, 32, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(32, 3, 3, 1, 1)
        )

    def forward(self, input):
        illumination = torch.sigmoid(self.illumination_net(input))
        reflectance = torch.sigmoid(self.reflectance_net(input))
        noise = torch.tanh(self.noise_net(input))

        return illumination, reflectance, noise

  我们可以对照上图左边的结构来理解代码。

  • illumination_net：  主要是负责对输入图像进行处理以获取光照信息，包括一系列卷积层和ReLU激活函数，最终输出一个通道数为1的图像，表示光照强度。

  • reflectance_net：  主要是负责提取输入图像的反射率信息，同样包括一系列卷积层和ReLU激活函数，最终输出一个通道数为3的图像，表示反射率在RGB通道上的分布。

  • noise_net：  主要是则用于估计输入图像的噪声信息，同样由一系列卷积层和ReLU激活函数组成，最终输出一个通道数为3的图像，表示噪声在RGB通道上的分布。

 最后，illumination_net和reflectance_net的输出经过sigmoid函数处理，而noise_net的输出则经过tanh函数处理。

 ☀️2.2 损失函数—loss_functions.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import conf

 #----------- 1.reconstruction_loss：计算重构损失---------------------------#
def reconstruction_loss(image, illumination, reflectance, noise):
    reconstructed_image = illumination*reflectance+noise
    return torch.norm(image-reconstructed_image, 1)

 #----------- 2.gradient: 计算输入图像在水平和垂直方向上的梯度--------------------#
def gradient(img):
    height = img.size(2)
    width = img.size(3)
    gradient_h = (img[:,:,2:,:]-img[:,:,:height-2,:]).abs()
    gradient_w = (img[:, :, :, 2:] - img[:, :, :, :width-2]).abs()
    gradient_h = F.pad(gradient_h, [0, 0, 1, 1], 'replicate')
    gradient_w = F.pad(gradient_w, [1, 1, 0, 0], 'replicate')
    gradient2_h = (img[:,:,4:,:]-img[:,:,:height-4,:]).abs()
    gradient2_w = (img[:, :, :, 4:] - img[:, :, :, :width-4]).abs()
    gradient2_h = F.pad(gradient2_h, [0, 0, 2, 2], 'replicate')
    gradient2_w = F.pad(gradient2_w, [2, 2, 0, 0], 'replicate')
    return gradient_h*gradient2_h, gradient_w*gradient2_w

 #----------- 3.normalize01: 将输入图像进行归一化到0到1的范围内---------------------#
def normalize01(img):
    minv = img.min()
    maxv = img.max()
    return (img-minv)/(maxv-minv)

 #----------- 4.gaussianblur3: 3通道的高斯模糊---------------------------#
def gaussianblur3(input):
    slice1 = F.conv2d(input[:,0,:,:].unsqueeze(1), weight=conf.gaussian_kernel, padding=conf.g_padding)
    slice2 = F.conv2d(input[:,1,:,:].unsqueeze(1), weight=conf.gaussian_kernel, padding=conf.g_padding)
    slice3 = F.conv2d(input[:,2,:,:].unsqueeze(1), weight=conf.gaussian_kernel, padding=conf.g_padding)
    x = torch.cat([slice1,slice2, slice3], dim=1)
    return x

 #----------- 5.illumination_smooth_loss: 计算光照平滑损失---------------------------#
def illumination_smooth_loss(image, illumination):
    gray_tensor = 0.299*image[0,0,:,:] + 0.587*image[0,1,:,:] + 0.114*image[0,2,:,:]
    max_rgb, _ = torch.max(image, 1)
    max_rgb = max_rgb.unsqueeze(1)
    gradient_gray_h, gradient_gray_w = gradient(gray_tensor.unsqueeze(0).unsqueeze(0))
    gradient_illu_h, gradient_illu_w = gradient(illumination)
    weight_h = 1/(F.conv2d(gradient_gray_h, weight=conf.gaussian_kernel, padding=conf.g_padding)+0.0001)
    weight_w = 1/(F.conv2d(gradient_gray_w, weight=conf.gaussian_kernel, padding=conf.g_padding)+0.0001)
    weight_h.detach()
    weight_w.detach()
    loss_h = weight_h * gradient_illu_h
    loss_w = weight_w * gradient_illu_w
    max_rgb.detach()
    return loss_h.sum() + loss_w.sum() + torch.norm(illumination-max_rgb, 1)

 #----------- 6.reflectance_smooth_loss:计算反射率平滑损失---------------------------#
def reflectance_smooth_loss(image, illumination, reflectance):
    gray_tensor = 0.299*image[0,0,:,:] + 0.587*image[0,1,:,:] + 0.114*image[0,2,:,:]
    gradient_gray_h, gradient_gray_w = gradient(gray_tensor.unsqueeze(0).unsqueeze(0))
    gradient_reflect_h, gradient_reflect_w = gradient(reflectance)
    weight = 1/(illumination*gradient_gray_h*gradient_gray_w+0.0001)
    weight = normalize01(weight)
    weight.detach()
    loss_h = weight * gradient_reflect_h
    loss_w = weight * gradient_reflect_w
    refrence_reflect = image/illumination
    refrence_reflect.detach()
    return loss_h.sum() + loss_w.sum() + conf.reffac*torch.norm(refrence_reflect - reflectance, 1)

 #----------- 7.noise_loss: 计算噪声损失---------------------------#
def noise_loss(image, illumination, reflectance, noise):
    weight_illu = illumination
    weight_illu.detach()
    loss = weight_illu*noise
    return torch.norm(loss, 2)
（1）重构损失——reconstruction_loss

图像的分解组件必须满足Robust Retinex的公式，将RGB三个通道中最大强度值S的初始值，在此基础上约束反射图和噪声。

（2）光照损失——illumination_smooth_loss

通过平滑的光照图可以增强暗区域的纹理细节，公式中x和y是水平和垂直方向，Wx和Wy是确保图像平滑的权重参数。

权重与梯度呈反比，梯度大的地方权重小，梯度小的地方权重大，因此将高斯滤波G放在分母，这里公式中的I是输入图像转换成的灰度图，Wy的计算方式和Wx的相同。

（3）反射损失——reflectance_smooth_loss

通过平滑反射图来得到噪声，本身并没有直接得到噪声的损失，只是通过对反射图做总变分约束来去噪。

（4）噪声损失——noise_loss

为了增加图像的清晰度增加了图像的对比度，与此同时，图像的噪声也被放大，出于以下两点限制噪声：

  1. 噪声的范围需要被限制。
  2. 噪声可以平滑的反射图限制。

  ☀️2.3 Retinex操作—pipline.py

import os
import numpy as np
import cv2
import torch
import torch.optim as optim
import torch.nn as nn
from PIL import Image
from torchvision import transforms
import torch.nn.init as init

from model.RRDNet import RRDNet
from loss.loss_functions import reconstruction_loss, illumination_smooth_loss, reflectance_smooth_loss, noise_loss, normalize01
import conf

 #----------- retinex图像增强---------------------------#
def pipline_retinex(net, img):
    img_tensor = transforms.ToTensor()(img)  # [c, h, w] #将输入图像转换为张量，并调整形状
    img_tensor = img_tensor.to(conf.device)
    img_tensor = img_tensor.unsqueeze(0)     # [1, c, h, w]

    optimizer = optim.Adam(net.parameters(), lr=conf.lr)

    # iterations：迭代优化过程
    for i in range(conf.iterations+1):
        # forward：通过网络前向传播得到光照、反射率和噪声图像。
        illumination, reflectance, noise = net(img_tensor)  # [1, c, h, w]
        # loss computing：计算总损失，并进行反向传播优化网络参数。
        loss_recons = reconstruction_loss(img_tensor, illumination, reflectance, noise)  # 重构损失
        loss_illu = illumination_smooth_loss(img_tensor, illumination) # 光照损失
        loss_reflect = reflectance_smooth_loss(img_tensor, illumination, reflectance) #反射损失
        loss_noise = noise_loss(img_tensor, illumination, reflectance, noise) # 噪声损失

        loss = loss_recons + conf.illu_factor*loss_illu + conf.reflect_factor*loss_reflect + conf.noise_factor*loss_noise

        # backward
        net.zero_grad()
        loss.backward()
        optimizer.step()

        # log：每隔 100 次迭代打印日志，显示重建损失、光照损失、反射率损失和噪声损失的数值。
        if i%100 == 0:
            print("iter:", i, '  reconstruction loss:', float(loss_recons.data), '  illumination loss:', float(loss_illu.data), '  reflectance loss:', float(loss_reflect.data), '  noise loss:', float(loss_noise.data))


    # adjustment：对增强后的图像进行调整
    adjust_illu = torch.pow(illumination, conf.gamma)
    res_image = adjust_illu*((img_tensor-noise)/illumination)# 对增强后的图像进行调整
    res_image = torch.clamp(res_image, min=0, max=1)# 对调整后的图像进行限幅操作，确保像素值在 0 到 1 之间。

    if conf.device != 'cpu':
        res_image = res_image.cpu()
        illumination = illumination.cpu()
        adjust_illu = adjust_illu.cpu()
        reflectance = reflectance.cpu()
        noise = noise.cpu()
    
    # 将处理后的张量转换为 PIL 图像
    res_img = transforms.ToPILImage()(res_image.squeeze(0))
    illum_img = transforms.ToPILImage()(illumination.squeeze(0))
    adjust_illu_img = transforms.ToPILImage()(adjust_illu.squeeze(0))
    reflect_img = transforms.ToPILImage()(reflectance.squeeze(0))
    noise_img = transforms.ToPILImage()(normalize01(noise.squeeze(0)))

    return res_img, illum_img, adjust_illu_img, reflect_img, noise_img


if __name__ == '__main__':

    # Init Model
    net = RRDNet()
    net = net.to(conf.device)

    # Test
    img = Image.open(conf.test_image_path)

    res_img, illum_img, adjust_illu_img, reflect_img, noise_img = pipline_retinex(net, img)
    res_img.save('./test/result.jpg')
    illum_img.save('./test/illumination.jpg')
    adjust_illu_img.save('./test/adjust_illumination.jpg')
    reflect_img.save('./test/reflectance.jpg')
    noise_img.save('./test/noise_map.jpg')

这段代码基本都注释了，就不再详细讲解了~

🚀三、RDDNet代码复现

☀️3.1 环境配置

  • Python 3
  • PyTorch >= 0.4.1
  • PIL >= 6.1.0
  • Opencv-python>=3.4

☀️3.2 运行过程

这个也是运行比较简单，配好环境就行 。不再过多叙述~

☀️3.3 运行效果

没错，你怎么知道我去看邓紫棋演唱会啦~ 

路人贾'ω'
关注
专栏目录
【代码笔记】RRDNet 网络
chaiky的博客
05-08 2628
论文笔记：低照度图像增强论文笔记1 代码地址：RRDNet_code 代码框架如图所示： 1、语法 import modname# 导入模块 import modname as name# 给导入的模块起别名 import library.modname as name# 导入包中的一个子模块，并将其命名为name from modname import name ...
代码解读——Retinex低光照图像增强（Deep Retinex Decomposition for Low-Light Enhancement）
热门推荐
sf_qw39的博客
03-07 1万+
今天带来一篇代码解读的文章，是2018年BMVC上的一篇暗光增强文章。个人觉得网络比较轻量并且能够取得还不错的效果。废话不多说，直接贴传送门： 文章地址：http://arxiv.org/abs/1808.04560 源码地址：https://github.com/weichen582/RetinexNet 文章基于Retinex理论，不懂的请戳这里：https://blog.csdn.ne...
ECCV 2020 论文大盘点-图像增强与图像恢复篇
我爱计算机视觉
10-15 3415
本文盘点 ECCV 2020 中底层图像处理方向相关的论文，包含：图像增强、图像恢复、去摩尔纹、去噪和质量评价论文，总计 29 篇，去摩尔纹 2 篇，去噪 10 篇，图像增强 7 篇，图...
【低照度图像增强系列（8）】URetinex-Net算法详解与代码实现（2022|CVPR）
最新发布
路人贾的博客
07-12 2423
【低照度图像增强系列（8）】URetinex-Net算法详解与代码实现（2022|CVPR）
论文笔记：ZERO-SHOT RESTORATION OF UNDEREXPOSED IMAGES VIA ROBUST RETINEX DECOMPOSITION
weixin_43535668的博客
10-26 746
论文名称ZERO-SHOT RESTORATION OF UNDEREXPOSED IMAGES VIA ROBUST RETINEX DECOMPOSITION 论文地址：Zero-Shot Restoration of Underexposed Images via Robust Retinex Decomposition | IEEE Conference Publication | IEEE Xplore 论文代码：https://aaaaangel.github.io/RRDNet-Homepag
发现暗夜的明眸：R2RNet——低光照图像增强的利器
gitblog_00038的博客
06-15 902
发现暗夜的明眸：R2RNet——低光照图像增强的利器 去发现同类优质开源项目:https://gitcode.com/ 在图像处理的世界里，点亮黑暗的每一个角落是一项挑战，而开源社区从未停止探索的脚步。今天，让我们一起揭开【R2RNet**】的神秘面纱，这款由Jiang Hai等研究人员打造的低光图像增强工具，它以“从真实低光到真实正常网络”（Real-low to Real-normal Net...
低照度图像增强论文阅读笔记——ZERO-SHOT RESTORATION OF UNDEREXPOSED IMAGES VIA ROBUST RETINEX DECOMPOSITION
weixin_50901244的博客
03-29 8206
ZERO-SHOT RESTORATION OF UNDEREXPOSED IMAGES VIA ROBUST RETINEX DECOMPOSITION1.摘要2.介绍3. RRDNet网络 1.摘要 低照度图像通常会出现严重的质量退化，如在黑暗中能见度差和潜在噪声。以往的低曝光图像复原方法大多忽略了噪声，在拉伸对比度的过程中对噪声进行放大。在增强低照度图像的同时，本文可以明确地预测噪声以达到去噪的目的。本文提出了一种新型的三分支卷积神经网络RRDNet(鲁棒Retinex分解网络的简称)，将输入图像分解
低照度图片增强——传统算法
Vaeeeee的博客
06-20 5572
光照因素一直是影响成像质量的一个关键因素，夜间等光照环境较差的条件下的图片往往细节丢失、分辨不清，信噪比低下。低照度图像增强是指通过一系列算法和技术，增强在低照度或弱光条件下拍摄的图像的可视化质量。本文主要介绍一些传统的低照度图像增强算法，给出具体的实现代码，便于测试各类图片的增强效果和比较各个算法的性能。以上几种经典的低照度图像增强算法，各自具有不同的优缺点，在实际应用中需要综合考虑图像的特点以及对图像质量的要求。
【论文笔记】—低照度图像增强—ZeroShot—RRDNet网络—2020-ICME
chaiky的博客
04-19 7131
题目：Zero-Shot Restoration Of Underexposed Images Via Robust Retinex Decomposition 时间：2020年 会议：2020-ICME 机构：同济大学 论文网址：https://sci-hub.se/10.1109/icme46284.2020.9102962 代码：https://github.com/aaaaangel/RRDNet 关键词：曝光不足的图像恢复、Retinex分解、zero-shot学习、去噪
【图像分割】基于ICM算法实现图像分割附matlab代码.zip
03-03
【图像分割】基于ICM算法实现图像分割附matlab代码.zip 这个压缩包文件包含了一系列与图像处理和机器学习相关的资源，特别是针对图像分割的Iterative Closest Point (ICM) 算法的实现。以下是这些资源背后的详细知识...
LIME低照度增强测试代码
01-03
LIME-low light/ night image enhancement 2017-tip 神作 效果惊人。包里含有完整的 测试代码； 直接可以使用作为对比
低照度图像增强（opencv和C++）
10-18
低照度图像增强，opencv3.2+vs2017，C++代码。 针对全局较暗的图像能有很好的增强效果。 代码注释的很详细，也比较简单。 提供了两张图片以供测试。 若有问题，可以找我。
【图像分割】基于迭代条件模式ICM算法实现图像分割matlab代码.zip
12-26
标题中的“【图像分割】基于迭代条件模式ICM算法实现图像分割matlab代码.zip”指出，这个压缩包包含了一个使用迭代条件模式（Iterated Conditional Modes, ICM）算法进行图像分割的MATLAB实现。图像分割是图像处理和...
ICM--MRF--matlab.rar_MRF 图像分割_MRF 算法_SAR MATLAB_SAR 图像_sar图像分割
07-14
MATLAB是一种广泛使用的编程环境，特别适合于数值计算和图像处理任务，因此这个压缩包可能是包含了一系列用MATLAB编写的代码或程序，用于实现MRF图像分割算法。 MRF是概率图模型的一种，常用于图像处理和计算机视觉...
【图像分割】基于迭代条件模式ICM算法实现图像分割matlab代码
qq_59747472的博客
12-26 1642
1 简介 基于MＲF的影像分割中，分割影像常用2个随机场来描述，一个是特征场(或称观测场) ，另一个是标记场。特征场以标记场为条件，用标记类别的似然分布函数描述特征数据或特征向量的分布;标记场则用先验知识描述标记场的局部相关性。影像分割需要将特征场中的数据与标记场中的先验知识概率建立有机的联系，这就需要将先验信息所期望的分割结果与真实观测数据的细节信息进行平衡，而贝叶斯决策理论能够很好地做到这一点。基于MＲF的影像分割通常用贝叶斯决策理论中的最大后验概率(maximum a poste...
【低照度图像增强系列（2）】Retinex（SSR/MSR/MSRCR）算法详解与代码实现
路人贾的博客
12-25 7581
低照度图像增强系列（2）Retinex算法详解与代码实现。包含SSR、MSR、MSRCR
探秘R2RNet：低光照图像增强的新里程碑
gitblog_00007的博客
06-15 992
探秘R2RNet：低光照图像增强的新里程碑 去发现同类优质开源项目:https://gitcode.com/ 在数字摄影领域，低光照条件下的图像质量一直是一个亟待解决的问题。【R2RNet】((https://www.sciencedirect.com/science/article/pii/S1047320322002322)[arxiv]) 是一个创新的解决方案，它提供了一个强大的框架来提升低...
图像融合亮度一致_基于深度学习的低光照图像增强方法总结（2020.09.15更新）...
weixin_39633165的博客
11-27 2142
之前在做光照对于高层视觉任务的影响的相关工作，看了不少基于深度学习的低光照增强（low-light enhancement）的文章，于是决定简单梳理一下。光照估计（illumination estimation）和低光照增强（low-light enhancement）的区别：光照估计是一个专门的底层视觉任务（例如[1,2,6]），它的输出结果可以被用到其它任务中，例如图像增强、图像恢复（处理色差...
低照度图像增强代码
jinx_ovo的博客
10-11 1333
低照度图像增强代码代码 代码 // An highlighted block #include <opencv2/core/core.hpp> #include <opencv2/imgproc/imgproc.hpp> #include <opencv2/highgui/highgui.hpp> #include <iostream> #include <stdio.h> //#include <cv.h> //#include&lt
用pytorch写添加了icm算法的ddpg代码
05-13
好的，以下是添加了ICM算法的DDPG代码： ```python import torch import torch.nn as nn import torch.nn.functional as F import numpy as np class ICM(nn.Module): def __init__(self, state_dim, action_dim, feature_dim=32, learning_rate=0.0001, beta=0.2): super(ICM, self).__init__() self.beta = beta self.inv_fc1 = nn.Linear(feature_dim*2, 256) self.inv_fc2 = nn.Linear(256, action_dim) self.forward_fc1 = nn.Linear(state_dim+action_dim, 256) self.forward_fc2 = nn.Linear(256, feature_dim) self.inverse_optimizer = torch.optim.Adam(self.inv_fc1.parameters(), lr=learning_rate) self.forward_optimizer = torch.optim.Adam(self.forward_fc1.parameters(), lr=learning_rate) def forward(self, state, next_state, action): state_action = torch.cat([state, action], 1) state_action = F.relu(self.forward_fc1(state_action)) next_state_feature = self.forward_fc2(state_action) state_feature = self.forward_fc2(F.relu(self.forward_fc1(torch.cat([state, torch.zeros_like(action)], 1)))) pred_action = self.inv_fc2(F.relu(self.inv_fc1(torch.cat([next_state_feature, state_feature], 1)))) return pred_action def compute_loss(self, state, next_state, action, pred_action): state_action = torch.cat([state, action], 1) state_action = F.relu(self.forward_fc1(state_action)) next_state_feature = self.forward_fc2(state_action) state_feature = self.forward_fc2(F.relu(self.forward_fc1(torch.cat([state, torch.zeros_like(action)], 1)))) inv_loss = F.mse_loss(pred_action, action) forward_loss = F.mse_loss(next_state_feature, state_feature.detach()) return self.beta * inv_loss + (1 - self.beta) * forward_loss def train_model(self, state, next_state, action): self.inverse_optimizer.zero_grad() self.forward_optimizer.zero_grad() pred_action = self(state, next_state, action) loss = self.compute_loss(state, next_state, action, pred_action) loss.backward() self.inverse_optimizer.step() self.forward_optimizer.step() class Critic(nn.Module): def __init__(self, state_dim, action_dim, hidden_dim1=128, hidden_dim2=128): super(Critic, self).__init__() self.fc1 = nn.Linear(state_dim + action_dim, hidden_dim1) self.fc2 = nn.Linear(hidden_dim1, hidden_dim2) self.fc3 = nn.Linear(hidden_dim2, 1) def forward(self, state, action): x = torch.cat([state, action], 1) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x class Actor(nn.Module): def __init__(self, state_dim, action_dim, max_action, hidden_dim1=128, hidden_dim2=128): super(Actor, self).__init__() self.fc1 = nn.Linear(state_dim, hidden_dim1) self.fc2 = nn.Linear(hidden_dim1, hidden_dim2) self.fc3 = nn.Linear(hidden_dim2, action_dim) self.max_action = max_action def forward(self, state): x = F.relu(self.fc1(state)) x = F.relu(self.fc2(x)) x = self.max_action * torch.tanh(self.fc3(x)) return x class DDPG: def __init__(self, state_dim, action_dim, max_action): self.actor = Actor(state_dim, action_dim, max_action) self.actor_target = Actor(state_dim, action_dim, max_action) self.actor_optimizer = torch.optim.Adam(self.actor.parameters()) self.critic = Critic(state_dim, action_dim) self.critic_target = Critic(state_dim, action_dim) self.critic_optimizer = torch.optim.Adam(self.critic.parameters()) self.icm = ICM(state_dim, action_dim) self.max_action = max_action def select_action(self, state): state = torch.FloatTensor(state.reshape(1, -1)) return self.actor(state).cpu().data.numpy().flatten() def update(self, replay_buffer, batch_size=100, discount=0.99, tau=0.005, icm_update=1): state, action, next_state, reward, done = replay_buffer.sample(batch_size) state = torch.FloatTensor(state) action = torch.FloatTensor(action) next_state = torch.FloatTensor(next_state) reward = torch.FloatTensor(reward.reshape(-1, 1)) done = torch.FloatTensor(done.reshape(-1, 1)) predicted_action = self.icm(state, next_state, action) # Q-function update target_Q = self.critic_target(next_state, self.actor_target(next_state)) target_Q = reward + ((1 - done) * discount * target_Q).detach() current_Q = self.critic(state, action) critic_loss = F.mse_loss(current_Q, target_Q) self.critic_optimizer.zero_grad() critic_loss.backward() self.critic_optimizer.step() # Actor update actor_loss = -self.critic(state, self.actor(state)).mean() self.actor_optimizer.zero_grad() actor_loss.backward() self.actor_optimizer.step() # ICM update if icm_update: icm_loss = self.icm.compute_loss(state, next_state, action, predicted_action) self.icm.train_model(state, next_state, action) # Target network update for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()): target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data) for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()): target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data) ``` 注意，这里的ICM算法是基于论文 "Curiosity-driven Exploration by Self-supervised Prediction" 实现的。在使用时，可以通过调节`icm_update`参数来指定是否更新ICM模型。如果将其设为0，则只运行DDPG算法。
评论 7
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

路人贾'ω'

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付：¥1
获取中
扫码支付

您的余额不足，请更换扫码支付或充值

打赏作者

实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值