bevy_rl 0.17.0

Build environments for reinforcement learning with bevy
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use bevy::{
    prelude::*,
    render::{
        render_asset::RenderAssets,
        render_resource::{
            Extent3d, TextureDescriptor, TextureDimension, TextureFormat, TextureUsages,
        },
        renderer::{RenderDevice, RenderQueue},
        texture::GpuImage,
    },
};

use wgpu::TexelCopyBufferLayout;
use wgpu::{BufferDescriptor, TexelCopyBufferInfo};

use crate::state;

fn texture_image_layout(desc: &TextureDescriptor<'_>) -> TexelCopyBufferLayout {
    let size = desc.size;

    let width = size.width * desc.format.block_dimensions().0;
    let height = size.width * desc.format.block_dimensions().1;

    TexelCopyBufferLayout {
        bytes_per_row: if size.height > 1 { Some(width) } else { None },
        rows_per_image: if size.depth_or_array_layers > 1 {
            Some(height)
        } else {
            None
        },
        ..Default::default()
    }
}

/// Copy a texture buffer from GPU to RAM and convert color space to RGBA.
/// It makes possible to export render results via API.
pub(crate) fn copy_from_gpu_to_ram<
    T: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe,
    P: 'static + Send + Sync + Clone + std::panic::RefUnwindSafe + serde::Serialize,
>(
    gpu_images: Res<RenderAssets<GpuImage>>,
    render_device: Res<RenderDevice>,
    render_queue: Res<RenderQueue>,
    ai_gym_state: Res<state::AIGymState<T, P>>,
) {
    let mut ai_gym_state_locked = ai_gym_state.lock().unwrap();
    if !ai_gym_state_locked.settings.render_to_buffer {
        return;
    }
    let ai_gym_settings = ai_gym_state_locked.settings.clone();

    let device = render_device.wgpu_device();
    let size = Extent3d {
        width: ai_gym_settings.width,
        height: ai_gym_settings.height,
        ..default()
    };

    ai_gym_state_locked.visual_observations = Vec::new();
    for (_, gp) in ai_gym_state_locked
        .render_image_handles
        .clone()
        .iter()
        .enumerate()
    {
        let render_gpu_image = gpu_images.get(gp).unwrap();
        let texture_width = size.width;
        let texture_height = size.height;

        let destination = device.create_buffer(&BufferDescriptor {
            label: None,
            size: (texture_width * texture_height * 4) as u64,
            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        let texture = render_gpu_image.texture.clone();

        let mut encoder =
            render_device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });

        let texture_extent = Extent3d {
            width: texture_width,
            height: texture_height,
            ..default()
        };

        encoder.copy_texture_to_buffer(
            texture.as_image_copy(),
            TexelCopyBufferInfo {
                buffer: &destination,
                layout: texture_image_layout(&TextureDescriptor {
                    label: None,
                    size,
                    dimension: TextureDimension::D2,
                    format: TextureFormat::Bgra8UnormSrgb,
                    mip_level_count: 1,
                    sample_count: 1,
                    usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
                    view_formats: &[TextureFormat::Bgra8UnormSrgb],
                }),
            },
            texture_extent,
        );

        render_queue.submit([encoder.finish()]);
        let buffer_slice = destination.slice(..);

        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
            let err = result.err();
            if err.is_some() {
                panic!("{}", err.unwrap().to_string());
            }
        });

        let _ = device.poll(wgpu::PollType::Wait);

        let data = buffer_slice.get_mapped_range();
        let result: Vec<u8> = bytemuck::cast_slice(&data).to_vec();

        drop(data);
        let mut rgba_image: image::RgbaImage =
            image::ImageBuffer::from_raw(texture_width, texture_height, result.clone()).unwrap();

        // fixing bgra to rgba
        convert_bgra_to_rgba(&mut rgba_image);

        ai_gym_state_locked
            .visual_observations
            .push(rgba_image.clone());

        destination.unmap();
    }
}

/// convert BRGA image to RGBA image
fn convert_bgra_to_rgba(image: &mut image::RgbaImage) {
    for pixel in image.pixels_mut() {
        pixel.0.swap(0, 2);
    }
}