thrust_rl/buffer/rollout/storage.rs
1//! Rollout buffer storage and data management
2//!
3//! This module handles the core storage functionality for rollout buffers,
4//! including data insertion, retrieval, and buffer management.
5//!
6//! The host-side storage layout (`Vec<f32>` / `Vec<i64>`) is deliberately
7//! backend-agnostic; tensor materialization happens via
8//! [`RolloutBatch::to_burn_tensors`] when the trainer pushes the batch
9//! to a device.
10
11use burn::tensor::{Int, Tensor as BurnTensor, TensorData, backend::Backend};
12
13/// Rollout buffer for storing trajectories
14///
15/// Stores trajectories collected from environment interactions and
16/// computes advantages using Generalized Advantage Estimation (GAE).
17///
18/// # Buffer Layout
19///
20/// The buffer uses a `[num_steps, num_envs]` layout where:
21/// - `num_steps`: Number of timesteps per rollout (typically 128-2048)
22/// - `num_envs`: Number of parallel environments
23///
24/// This layout provides good cache locality for forward passes and
25/// efficient computation of advantages.
26#[derive(Debug, Clone)]
27pub struct RolloutBuffer {
28 /// Number of steps per rollout
29 num_steps: usize,
30
31 /// Number of parallel environments
32 num_envs: usize,
33
34 /// Dimensionality of observations
35 obs_dim: usize,
36
37 /// Observations [num_steps, num_envs, obs_dim]
38 observations: Vec<Vec<Vec<f32>>>,
39
40 /// Actions taken [num_steps, num_envs]
41 actions: Vec<Vec<i64>>,
42
43 /// Rewards received [num_steps, num_envs]
44 rewards: Vec<Vec<f32>>,
45
46 /// Value estimates [num_steps, num_envs]
47 values: Vec<Vec<f32>>,
48
49 /// Log probabilities [num_steps, num_envs]
50 log_probs: Vec<Vec<f32>>,
51
52 /// Episode termination flags [num_steps, num_envs]
53 terminated: Vec<Vec<bool>>,
54
55 /// Episode truncation flags [num_steps, num_envs]
56 truncated: Vec<Vec<bool>>,
57
58 /// Computed advantages [num_steps, num_envs]
59 advantages: Vec<Vec<f32>>,
60
61 /// Computed returns [num_steps, num_envs]
62 returns: Vec<Vec<f32>>,
63}
64
65impl RolloutBuffer {
66 /// Create a new rollout buffer
67 ///
68 /// # Arguments
69 ///
70 /// * `num_steps` - Number of timesteps per rollout
71 /// * `num_envs` - Number of parallel environments
72 /// * `obs_dim` - Dimensionality of observations
73 pub fn new(num_steps: usize, num_envs: usize, obs_dim: usize) -> Self {
74 // Pre-allocate all buffers
75 let observations = vec![vec![vec![0.0; obs_dim]; num_envs]; num_steps];
76 let actions = vec![vec![0; num_envs]; num_steps];
77 let rewards = vec![vec![0.0; num_envs]; num_steps];
78 let values = vec![vec![0.0; num_envs]; num_steps];
79 let log_probs = vec![vec![0.0; num_envs]; num_steps];
80 let terminated = vec![vec![false; num_envs]; num_steps];
81 let truncated = vec![vec![false; num_envs]; num_steps];
82 let advantages = vec![vec![0.0; num_envs]; num_steps];
83 let returns = vec![vec![0.0; num_envs]; num_steps];
84
85 Self {
86 num_steps,
87 num_envs,
88 obs_dim,
89 observations,
90 actions,
91 rewards,
92 values,
93 log_probs,
94 terminated,
95 truncated,
96 advantages,
97 returns,
98 }
99 }
100
101 /// Add a transition to the buffer
102 ///
103 /// # Arguments
104 ///
105 /// * `step` - Timestep within the rollout (0 to num_steps-1)
106 /// * `env_id` - Environment ID (0 to num_envs-1)
107 /// * `observation` - Current observation
108 /// * `action` - Action taken
109 /// * `reward` - Reward received
110 /// * `value` - Value estimate for current state
111 /// * `log_prob` - Log probability of the action
112 /// * `terminated` - Whether the episode terminated
113 /// * `truncated` - Whether the episode was truncated
114 // Each argument is a distinct transition field; bundling them into a struct
115 // would add boilerplate at every call site without improving clarity.
116 #[allow(clippy::too_many_arguments)]
117 pub fn add(
118 &mut self,
119 step: usize,
120 env_id: usize,
121 observation: &[f32],
122 action: i64,
123 reward: f32,
124 value: f32,
125 log_prob: f32,
126 terminated: bool,
127 truncated: bool,
128 ) {
129 debug_assert!(step < self.num_steps, "step {} >= num_steps {}", step, self.num_steps);
130 debug_assert!(env_id < self.num_envs, "env_id {} >= num_envs {}", env_id, self.num_envs);
131 debug_assert_eq!(observation.len(), self.obs_dim, "observation dimension mismatch");
132
133 self.observations[step][env_id].copy_from_slice(observation);
134 self.actions[step][env_id] = action;
135 self.rewards[step][env_id] = reward;
136 self.values[step][env_id] = value;
137 self.log_probs[step][env_id] = log_prob;
138 self.terminated[step][env_id] = terminated;
139 self.truncated[step][env_id] = truncated;
140 }
141
142 /// Reset the buffer for a new rollout
143 pub fn reset(&mut self) {
144 // Clear computed advantages and returns
145 for step in 0..self.num_steps {
146 for env in 0..self.num_envs {
147 self.advantages[step][env] = 0.0;
148 self.returns[step][env] = 0.0;
149 }
150 }
151 }
152
153 /// Get buffer shape (num_steps, num_envs, obs_dim)
154 pub fn shape(&self) -> (usize, usize, usize) {
155 (self.num_steps, self.num_envs, self.obs_dim)
156 }
157
158 /// Get total number of transitions in buffer
159 pub fn len(&self) -> usize {
160 self.num_steps * self.num_envs
161 }
162
163 /// Check if buffer is empty
164 pub fn is_empty(&self) -> bool {
165 self.len() == 0
166 }
167
168 /// Get observations tensor shape for neural network input
169 pub fn obs_shape(&self) -> (usize, usize) {
170 (self.num_steps * self.num_envs, self.obs_dim)
171 }
172
173 // ---- Getters for raw data access ----
174 //
175 // All getters return a slice indexed as `[step][env]` (outer = time,
176 // inner = parallel env). GAE, log-prob masking, and minibatch sampling
177 // all rely on this layout; do not transpose at the call site without
178 // adjusting [`compute_advantages`](super::gae::compute_advantages) accordingly.
179
180 /// Borrow the per-step observations, indexed `[step][env]` and
181 /// then by observation dimension (final inner `Vec<f32>` is one
182 /// observation vector of length `obs_dim`).
183 pub fn observations(&self) -> &[Vec<Vec<f32>>] {
184 &self.observations
185 }
186 /// Borrow the per-step discrete actions, indexed `[step][env]`.
187 pub fn actions(&self) -> &[Vec<i64>] {
188 &self.actions
189 }
190 /// Borrow the per-step rewards (raw, un-discounted), indexed
191 /// `[step][env]`.
192 pub fn rewards(&self) -> &[Vec<f32>] {
193 &self.rewards
194 }
195 /// Borrow the per-step bootstrap value estimates `V(s_t)` produced
196 /// by the policy at collection time, indexed `[step][env]`. Used by
197 /// [`compute_advantages`](super::gae::compute_advantages) as the value
198 /// baseline.
199 pub fn values(&self) -> &[Vec<f32>] {
200 &self.values
201 }
202 /// Borrow the per-step action log-probabilities under the
203 /// behavior policy at collection time, indexed `[step][env]`. PPO
204 /// uses these as the denominator of the importance-sampling ratio.
205 pub fn log_probs(&self) -> &[Vec<f32>] {
206 &self.log_probs
207 }
208 /// Borrow the per-step terminal-state flags (true on episode end),
209 /// indexed `[step][env]`. GAE zeroes the bootstrap across terminal
210 /// transitions but keeps the realized reward.
211 pub fn terminated(&self) -> &[Vec<bool>] {
212 &self.terminated
213 }
214 /// Borrow the per-step truncation flags (true on time-limit / external
215 /// reset that is not a terminal state), indexed `[step][env]`. GAE
216 /// retains the bootstrap value across truncated transitions because the
217 /// trajectory is still "alive" in the value-function sense.
218 pub fn truncated(&self) -> &[Vec<bool>] {
219 &self.truncated
220 }
221 /// Borrow the per-step GAE advantages computed by
222 /// [`compute_advantages`](super::gae::compute_advantages), indexed
223 /// `[step][env]`. Zero-initialized until GAE has run.
224 pub fn advantages(&self) -> &[Vec<f32>] {
225 &self.advantages
226 }
227 /// Borrow the per-step value-function targets (advantages + values),
228 /// indexed `[step][env]`. Zero-initialized until GAE has run.
229 pub fn returns(&self) -> &[Vec<f32>] {
230 &self.returns
231 }
232
233 // ---- Mutable getters for advantage/return computation ----
234
235 /// Mutable view of the advantages buffer, indexed `[step][env]`.
236 /// Used when in-place normalizing advantages or applying a custom
237 /// advantage estimator that does not touch `returns`; use
238 /// [`Self::advantages_and_returns_mut`] when both must be borrowed
239 /// at the same time (e.g. inside
240 /// [`compute_advantages`](super::gae::compute_advantages)).
241 pub fn advantages_mut(&mut self) -> &mut [Vec<f32>] {
242 &mut self.advantages
243 }
244 /// Mutable view of the returns / value-target buffer, indexed
245 /// `[step][env]`. Use [`Self::advantages_and_returns_mut`] instead
246 /// when both must be borrowed at the same time.
247 pub fn returns_mut(&mut self) -> &mut [Vec<f32>] {
248 &mut self.returns
249 }
250
251 /// Get mutable references to both advantages and returns
252 /// This is needed to avoid double mutable borrow in GAE computation
253 pub fn advantages_and_returns_mut(&mut self) -> (&mut [Vec<f32>], &mut [Vec<f32>]) {
254 (&mut self.advantages, &mut self.returns)
255 }
256}
257
258/// Batch of rollout data for training
259///
260/// Contains flattened tensors suitable for neural network training.
261/// All arrays have shape `[batch_size]`.
262#[derive(Debug, Clone)]
263pub struct RolloutBatch {
264 /// Flattened observations [batch_size, obs_dim]
265 pub observations: Vec<f32>,
266
267 /// Actions taken `[batch_size]`
268 pub actions: Vec<i64>,
269
270 /// Old log probabilities `[batch_size]`
271 pub old_log_probs: Vec<f32>,
272
273 /// Old value estimates `[batch_size]`
274 pub old_values: Vec<f32>,
275
276 /// Computed advantages `[batch_size]`
277 pub advantages: Vec<f32>,
278
279 /// Computed returns `[batch_size]`
280 pub returns: Vec<f32>,
281}
282
283impl RolloutBatch {
284 /// Create a new batch from rollout buffer
285 ///
286 /// Iterates the full `[num_steps, num_envs]` capacity. If the buffer
287 /// was only partially filled, the unwritten tail surfaces as
288 /// zero-initialized rows. Use [`Self::from_buffer_partial`] (or
289 /// [`super::RolloutBuffer::get_filled_batch`]) when the caller
290 /// knows the fill count.
291 pub fn from_buffer(buffer: &RolloutBuffer) -> Self {
292 Self::from_buffer_partial(buffer, buffer.num_steps)
293 }
294
295 /// Create a new batch from the first `valid_steps` rows of the buffer.
296 ///
297 /// Rows in `valid_steps..num_steps` (the unfilled tail of a partial
298 /// rollout) are skipped, preventing zero-padded rows from
299 /// contaminating PPO gradients.
300 ///
301 /// # Panics
302 /// Panics if `valid_steps > buffer.num_steps`.
303 pub fn from_buffer_partial(buffer: &RolloutBuffer, valid_steps: usize) -> Self {
304 assert!(
305 valid_steps <= buffer.num_steps,
306 "valid_steps ({}) must not exceed buffer.num_steps ({})",
307 valid_steps,
308 buffer.num_steps
309 );
310
311 let batch_size = valid_steps * buffer.num_envs;
312 let obs_size = batch_size * buffer.obs_dim;
313
314 let mut observations = Vec::with_capacity(obs_size);
315 let mut actions = Vec::with_capacity(batch_size);
316 let mut old_log_probs = Vec::with_capacity(batch_size);
317 let mut old_values = Vec::with_capacity(batch_size);
318 let mut advantages = Vec::with_capacity(batch_size);
319 let mut returns = Vec::with_capacity(batch_size);
320
321 // Flatten the filled prefix into 1D arrays
322 for step in 0..valid_steps {
323 for env in 0..buffer.num_envs {
324 observations.extend_from_slice(&buffer.observations[step][env]);
325 actions.push(buffer.actions[step][env]);
326 old_log_probs.push(buffer.log_probs[step][env]);
327 old_values.push(buffer.values[step][env]);
328 advantages.push(buffer.advantages[step][env]);
329 returns.push(buffer.returns[step][env]);
330 }
331 }
332
333 Self { observations, actions, old_log_probs, old_values, advantages, returns }
334 }
335
336 /// Get batch size
337 pub fn len(&self) -> usize {
338 self.actions.len()
339 }
340
341 /// Check if batch is empty
342 pub fn is_empty(&self) -> bool {
343 self.len() == 0
344 }
345
346 /// Get the observation shape as (batch_size, obs_dim)
347 pub fn obs_shape(&self) -> (usize, usize) {
348 let batch_size = self.len();
349 let obs_dim = self.observations.len().checked_div(batch_size).unwrap_or(0);
350 (batch_size, obs_dim)
351 }
352
353 /// Construct the full set of training tensors from this batch in a
354 /// Construct the full set of training tensors as Burn tensors on
355 /// `device`.
356 ///
357 /// Returns a named [`RolloutBurnTensors`] bundle so trainers can
358 /// pattern-match named fields rather than positional tuple elements.
359 ///
360 /// Shapes (all on `device`):
361 /// - `observations`: `[batch, obs_dim]`, `f32`
362 /// - `actions`: `[batch]`, `i64` (Burn `Int` kind)
363 /// - `old_log_probs`: `[batch]`, `f32`
364 /// - `old_values`: `[batch]`, `f32`
365 /// - `advantages`: `[batch]`, `f32`
366 /// - `returns`: `[batch]`, `f32`
367 ///
368 /// Empty batches still produce well-formed `[0, obs_dim]` /
369 /// `[0]` tensors. The `obs_dim` is derived from
370 /// [`Self::obs_shape`] (which returns `0` for an empty batch), so
371 /// the observation tensor for the empty case is shaped `[0, 0]`.
372 pub fn to_burn_tensors<B: Backend>(&self, device: &B::Device) -> RolloutBurnTensors<B> {
373 let (batch_size, obs_dim) = self.obs_shape();
374
375 // Construct the rank-2 observation tensor directly rather than
376 // building a rank-1 tensor and reshaping. The reshape path hits a
377 // panic deep in cubecl-zspace when both dims are zero (empty
378 // batch), and direct rank-2 construction sidesteps that edge case.
379 let observations = BurnTensor::<B, 2>::from_data(
380 TensorData::new(self.observations.clone(), [batch_size, obs_dim]),
381 device,
382 );
383 let actions = BurnTensor::<B, 1, Int>::from_data(
384 TensorData::new(self.actions.clone(), [batch_size]),
385 device,
386 );
387 let old_log_probs = BurnTensor::<B, 1>::from_data(
388 TensorData::new(self.old_log_probs.clone(), [batch_size]),
389 device,
390 );
391 let old_values = BurnTensor::<B, 1>::from_data(
392 TensorData::new(self.old_values.clone(), [batch_size]),
393 device,
394 );
395 let advantages = BurnTensor::<B, 1>::from_data(
396 TensorData::new(self.advantages.clone(), [batch_size]),
397 device,
398 );
399 let returns = BurnTensor::<B, 1>::from_data(
400 TensorData::new(self.returns.clone(), [batch_size]),
401 device,
402 );
403
404 RolloutBurnTensors { observations, actions, old_log_probs, old_values, advantages, returns }
405 }
406}
407
408/// Bundle of Burn tensors produced by [`RolloutBatch::to_burn_tensors`].
409///
410/// Fields are in the order PPO trainers consume them: policy/value
411/// inputs first (observations, actions), then the old policy outputs
412/// (log-probs and values used for the importance ratio and value clip),
413/// then the GAE outputs (advantages and returns). Generic over the
414/// backend `B` so the same trainer surface works for CPU (`NdArray`),
415/// GPU (`Wgpu`, `Cuda`), and `Autodiff<_>` wrappers.
416#[derive(Debug)]
417pub struct RolloutBurnTensors<B: Backend> {
418 /// Observations, shape `[batch_size, obs_dim]`, dtype `f32`.
419 pub observations: BurnTensor<B, 2>,
420
421 /// Discrete actions, shape `[batch_size]`, dtype `i64`.
422 pub actions: BurnTensor<B, 1, Int>,
423
424 /// Behavior-policy log-probabilities, shape `[batch_size]`,
425 /// dtype `f32`.
426 pub old_log_probs: BurnTensor<B, 1>,
427
428 /// Behavior-policy value estimates `V(s_t)`, shape `[batch_size]`,
429 /// dtype `f32`.
430 pub old_values: BurnTensor<B, 1>,
431
432 /// GAE advantages, shape `[batch_size]`, dtype `f32`.
433 pub advantages: BurnTensor<B, 1>,
434
435 /// Value-function targets (advantages + values), shape
436 /// `[batch_size]`, dtype `f32`.
437 pub returns: BurnTensor<B, 1>,
438}
439
440#[cfg(test)]
441mod tests {
442 use super::*;
443
444 #[test]
445 fn test_rollout_buffer_creation() {
446 let buffer = RolloutBuffer::new(10, 2, 4);
447
448 assert_eq!(buffer.shape(), (10, 2, 4));
449 assert_eq!(buffer.len(), 20); // 10 steps * 2 envs
450 assert!(!buffer.is_empty());
451 }
452
453 #[test]
454 fn test_rollout_buffer_add_and_reset() {
455 let mut buffer = RolloutBuffer::new(5, 1, 2);
456
457 // Add some data
458 buffer.add(0, 0, &[1.0, 2.0], 1, 1.5, 0.8, -0.2, false, false);
459 buffer.add(1, 0, &[2.0, 3.0], 0, 2.0, 1.2, -0.1, false, false);
460
461 // Check data was stored
462 assert_eq!(buffer.actions()[0][0], 1);
463 assert_eq!(buffer.rewards()[0][0], 1.5);
464 assert_eq!(buffer.observations()[0][0], vec![1.0, 2.0]);
465
466 // Reset and check advantages/returns are cleared
467 buffer.reset();
468 assert_eq!(buffer.advantages()[0][0], 0.0);
469 assert_eq!(buffer.returns()[0][0], 0.0);
470 }
471
472 #[test]
473 fn test_rollout_batch_from_buffer() {
474 let mut buffer = RolloutBuffer::new(2, 1, 2);
475
476 // Add test data
477 buffer.add(0, 0, &[1.0, 2.0], 1, 1.5, 0.8, -0.2, false, false);
478 buffer.add(1, 0, &[2.0, 3.0], 0, 2.0, 1.2, -0.1, false, false);
479
480 // Set some advantages and returns
481 buffer.advantages_mut()[0][0] = 0.5;
482 buffer.returns_mut()[0][0] = 1.3;
483 buffer.advantages_mut()[1][0] = 0.8;
484 buffer.returns_mut()[1][0] = 2.0;
485
486 let batch = RolloutBatch::from_buffer(&buffer);
487
488 assert_eq!(batch.len(), 2);
489 assert_eq!(batch.actions, vec![1, 0]);
490 assert_eq!(batch.advantages, vec![0.5, 0.8]);
491 assert_eq!(batch.returns, vec![1.3, 2.0]);
492 assert_eq!(batch.observations, vec![1.0, 2.0, 2.0, 3.0]);
493 }
494
495 #[test]
496 fn test_rollout_batch_from_buffer_partial_skips_unfilled_tail() {
497 // 4-step buffer, single env, only the first 2 rows filled.
498 let mut buffer = RolloutBuffer::new(4, 1, 2);
499
500 buffer.add(0, 0, &[1.0, 2.0], 1, 1.5, 0.8, -0.2, false, false);
501 buffer.add(1, 0, &[2.0, 3.0], 0, 2.0, 1.2, -0.1, false, false);
502 buffer.advantages_mut()[0][0] = 0.5;
503 buffer.returns_mut()[0][0] = 1.3;
504 buffer.advantages_mut()[1][0] = 0.8;
505 buffer.returns_mut()[1][0] = 2.0;
506
507 let batch = RolloutBatch::from_buffer_partial(&buffer, 2);
508
509 // Batch has exactly 2 rows — the zero-initialized tail (rows 2-3)
510 // is skipped.
511 assert_eq!(batch.len(), 2);
512 assert_eq!(batch.actions, vec![1, 0]);
513 assert_eq!(batch.old_values, vec![0.8, 1.2]);
514 assert_eq!(batch.old_log_probs, vec![-0.2, -0.1]);
515 assert_eq!(batch.advantages, vec![0.5, 0.8]);
516 assert_eq!(batch.returns, vec![1.3, 2.0]);
517 assert_eq!(batch.observations, vec![1.0, 2.0, 2.0, 3.0]);
518
519 // `from_buffer` (full capacity) still emits 4 rows; the 2-row
520 // partial batch is a strict subset.
521 let full_batch = RolloutBatch::from_buffer(&buffer);
522 assert_eq!(full_batch.len(), 4);
523 }
524
525 #[test]
526 fn test_rollout_batch_from_buffer_partial_zero_valid() {
527 let buffer = RolloutBuffer::new(4, 2, 3);
528 let batch = RolloutBatch::from_buffer_partial(&buffer, 0);
529 assert!(batch.is_empty());
530 assert_eq!(batch.actions.len(), 0);
531 assert_eq!(batch.observations.len(), 0);
532 }
533
534 #[test]
535 #[should_panic(expected = "valid_steps")]
536 fn test_rollout_batch_from_buffer_partial_panics_on_overflow() {
537 let buffer = RolloutBuffer::new(4, 1, 2);
538 let _ = RolloutBatch::from_buffer_partial(&buffer, 5);
539 }
540
541 mod burn_tests {
542 use burn::backend::NdArray;
543
544 use super::*;
545
546 type B = NdArray<f32>;
547
548 #[test]
549 fn test_to_burn_tensors_matches_inline_construction() {
550 // Mirrors the tch round-trip test: shapes + element-wise equality
551 // against the source `Vec`s the batch was built from.
552 let batch = RolloutBatch {
553 observations: vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
554 actions: vec![0, 1, 0],
555 old_log_probs: vec![-0.1, -0.2, -0.3],
556 old_values: vec![0.5, 0.7, 0.9],
557 advantages: vec![0.1, -0.2, 0.3],
558 returns: vec![0.6, 0.5, 1.2],
559 };
560
561 let device = crate::utils::cuda::default_burn_device::<B>();
562 let t = batch.to_burn_tensors::<B>(&device);
563
564 // Shapes.
565 assert_eq!(t.observations.dims(), [3, 2]);
566 assert_eq!(t.actions.dims(), [3]);
567 assert_eq!(t.old_log_probs.dims(), [3]);
568 assert_eq!(t.old_values.dims(), [3]);
569 assert_eq!(t.advantages.dims(), [3]);
570 assert_eq!(t.returns.dims(), [3]);
571
572 // Round-trip values. Observations are row-major flatten, so
573 // copying the `[3, 2]` tensor back to a `Vec<f32>` must equal
574 // the original 1-D buffer.
575 let obs_flat: Vec<f32> = t.observations.into_data().to_vec().unwrap();
576 assert_eq!(obs_flat, batch.observations);
577 let acts: Vec<i64> = t.actions.into_data().to_vec().unwrap();
578 assert_eq!(acts, batch.actions);
579 let lp: Vec<f32> = t.old_log_probs.into_data().to_vec().unwrap();
580 assert_eq!(lp, batch.old_log_probs);
581 let v: Vec<f32> = t.old_values.into_data().to_vec().unwrap();
582 assert_eq!(v, batch.old_values);
583 let adv: Vec<f32> = t.advantages.into_data().to_vec().unwrap();
584 assert_eq!(adv, batch.advantages);
585 let ret: Vec<f32> = t.returns.into_data().to_vec().unwrap();
586 assert_eq!(ret, batch.returns);
587 }
588
589 #[test]
590 fn test_to_burn_tensors_empty_batch() {
591 // Empty batch → `[0, 0]` observation tensor and `[0]`
592 // scalar-per-row tensors, matching the tch path's edge case.
593 let batch = RolloutBatch {
594 observations: vec![],
595 actions: vec![],
596 old_log_probs: vec![],
597 old_values: vec![],
598 advantages: vec![],
599 returns: vec![],
600 };
601
602 let device = crate::utils::cuda::default_burn_device::<B>();
603 let t = batch.to_burn_tensors::<B>(&device);
604
605 assert_eq!(t.observations.dims(), [0, 0]);
606 assert_eq!(t.actions.dims(), [0]);
607 assert_eq!(t.old_log_probs.dims(), [0]);
608 assert_eq!(t.old_values.dims(), [0]);
609 assert_eq!(t.advantages.dims(), [0]);
610 assert_eq!(t.returns.dims(), [0]);
611 }
612 }
613
614 #[test]
615 fn test_rollout_batch_properties() {
616 let batch = RolloutBatch {
617 observations: vec![1.0, 2.0, 3.0, 4.0],
618 actions: vec![0, 1],
619 old_log_probs: vec![-0.1, -0.2],
620 old_values: vec![0.5, 0.8],
621 advantages: vec![0.3, 0.6],
622 returns: vec![1.0, 1.5],
623 };
624
625 assert_eq!(batch.len(), 2);
626 assert_eq!(batch.obs_shape(), (2, 2)); // 2 samples, 2 obs dims each
627 assert!(!batch.is_empty());
628 }
629}