relayrl_types

Core data types and encoding/decoding utilities for the RelayRL framework.

Features

RelayRLAction: Serializable action container (obs, act, mask, reward, data, done) with UUID agent tracking
RelayRLTrajectory: In-memory trajectory buffer with metadata and provenance tracking
Records: Records for converting trajectories to and from CSV and Arrow files
Burn backend support: Compatible with both burn-ndarray (CPU) and burn-tch (GPU) backends
Codec pipeline: Compression, encryption, integrity verification, and chunking
Utilities: Metadata tracking, quantization, and network transport optimizations

Feature Flags

The crate's actual default set is ["ndarray-backend", "onnx-model"] — the CPU Burn backend plus ONNX inference. Codec features (compression, encryption, integrity, metadata, quantization, zerocopy) are not enabled by default; the codec examples below require opting into them explicitly (individually, or via a bundle like codec-full).

[features]
default = ["ndarray-backend", "onnx-model"]

# Backend selection
tch-backend = ["burn-tch", "half"]          # GPU (LibTorch) backend
ndarray-backend = ["burn-ndarray", "half"]  # CPU backend

# Inference model formats
inference-models = ["tch-model", "onnx-model"]  # All inference models
tch-model = ["tch", "tokio", "tempfile"]        # LibTorch inference
onnx-model = ["ort", "tokio", "tempfile", "ndarray"]  # ONNX inference

# Codec pipeline (opt-in)
compression = ["lz4_flex", "zstd", "bincode"]  # LZ4/Zstd compression
encryption = ["chacha20poly1305", "bincode"]   # ChaCha20-Poly1305 AEAD
integrity = ["blake3"]                          # BLAKE3 checksums
metadata = ["bincode"]                          # `encode`/`decode`/`to_bytes`/`from_bytes`
quantization = ["half"]                         # FP16/BF16 quantization
zerocopy = ["bytes"]                            # Zerocopy data conversions

# Convenience bundles
codec-basic = ["compression", "integrity", "zerocopy"]
codec-secure = ["codec-basic", "encryption"]
codec-full = ["codec-secure", "metadata", "quantization"]

To run the codec examples below, add relayrl_types with features = ["codec-full"] (or the specific features each example needs) in addition to the default backend.

Quick Start

Basic Usage

use relayrl_types::prelude::action::RelayRLAction;
use relayrl_types::prelude::trajectory::RelayRLTrajectory;
use relayrl_types::prelude::tensor::relayrl::{DType, NdArrayDType, DeviceType, TensorData};
use relayrl_types::data::tensor::ConversionBurnTensor;
use uuid::Uuid;
use std::sync::Arc;
use burn_tensor::Tensor;
use burn_ndarray::NdArray; // enable feature: ndarray-backend

// Create a Burn tensor (NdArray backend) and store as RelayRL TensorData
let device = DeviceType::Cpu;

// 1) Burn → RelayRL: Convert any Burn tensor into TensorData with a target dtype/backend.
//    `ConversionBurnTensor` wraps the tensor in an `Arc` (it may be shared across conversions).
let obs_burn = Tensor::<NdArray, 1>::from_floats([1.0, 2.0, 3.0, 4.0], &burn_tensor::Device::Cpu);
let obs_td: TensorData = ConversionBurnTensor {
    inner: Arc::new(obs_burn),
    conversion_dtype: DType::NdArray(NdArrayDType::F32),
}.try_into()?;

let act_burn = Tensor::<NdArray, 1>::from_floats([0.5, -0.3], &burn_tensor::Device::Cpu);
let act_td: TensorData = ConversionBurnTensor {
    inner: Arc::new(act_burn),
    conversion_dtype: DType::NdArray(NdArrayDType::F32),
}.try_into()?;

// 2) RelayRL → Burn: Build Burn tensors from stored TensorData with a chosen backend/device
//    Specify the backend type parameter; device is provided via DeviceType
let obs_tensor_any = RelayRLAction::to_tensor::<NdArray>(&obs_td, &device)?; // Box<dyn Any>
let act_tensor_any = RelayRLAction::to_tensor::<NdArray>(&act_td, &device)?;

// 3) Create an action with tensors
let action = RelayRLAction::new(
    Some(obs_td),                 // observation TensorData
    Some(act_td),                 // action TensorData
    None,                         // mask
    1.5,                          // reward
    false,                        // done
    None,                         // auxiliary data
    Some(Uuid::new_v4()),        // agent_id
);

// 4) Work with a trajectory
let mut trajectory = RelayRLTrajectory::with_agent_id(1000, Uuid::new_v4());
trajectory.add_action(action);

println!("Total reward: {}", trajectory.total_reward());
println!("Length: {}", trajectory.len());

// Minimal action without tensors
trajectory.add_action(RelayRLAction::minimal(1.0, false));

Codec Functionality

1. Simple Serialization

use relayrl_types::prelude::action::RelayRLAction;

let action = RelayRLAction::minimal(1.0, false);

// Simple serialization (requires the "metadata" feature).
// `from_bytes` returns `(Self, usize)`: the decoded value and the number of bytes consumed.
let bytes = action.to_bytes()?;
let (decoded, _consumed) = RelayRLAction::from_bytes(&bytes)?;

assert_eq!(decoded.get_rew(), 1.0);

2. Compression

use relayrl_types::prelude::trajectory::RelayRLTrajectory;
use relayrl_types::prelude::action::CodecConfig;
use relayrl_types::prelude::codec::CompressionScheme;

let trajectory = RelayRLTrajectory::new(100);
// ... add actions ...

// Configure codec with LZ4 compression (fast). Requires "compression", "integrity",
// and "metadata" features (e.g. via the "codec-full" bundle).
let config = CodecConfig {
    compression: Some(CompressionScheme::Lz4),
    encryption_key: None,
    verify_integrity: true,
    include_metadata: true,
};

// Encode with compression
let encoded = trajectory.encode(&config)?;
println!("Compressed from {} to {} bytes", 
    encoded.original_size, 
    encoded.data.len()
);

// Decode. `RelayRLTrajectory::decode` returns `(Self, usize)`.
let (decoded, _consumed) = RelayRLTrajectory::decode(&encoded, &config)?;

3. Compression + Encryption

use relayrl_types::prelude::action::{RelayRLAction, CodecConfig};
use relayrl_types::prelude::codec::CompressionScheme;
use relayrl_types::data::utilities::encrypt::generate_key;

let action = RelayRLAction::minimal(2.5, true);

// Generate encryption key (requires the "encryption" feature)
let key = generate_key();

// Configure codec with compression AND encryption.
// Requires "compression", "encryption", "integrity", and "metadata" features.
let config = CodecConfig {
    compression: Some(CompressionScheme::Zstd(3)),  // Zstd level 3
    encryption_key: Some(key),
    verify_integrity: true,
    include_metadata: true,
};

// Encode (compressed + encrypted)
let encoded = action.encode(&config)?;

// Decode (must use same key!). `RelayRLAction::decode` returns `Self` directly.
let decoded = RelayRLAction::decode(&encoded, &config)?;
assert_eq!(decoded.get_rew(), 2.5);

4. Full Pipeline with Integrity Verification

use relayrl_types::prelude::action::{RelayRLAction, CodecConfig};
use relayrl_types::prelude::trajectory::RelayRLTrajectory;
use relayrl_types::prelude::codec::CompressionScheme;
use relayrl_types::data::utilities::encrypt::generate_key;

let mut trajectory = RelayRLTrajectory::new(100);
for i in 0..50 {
    trajectory.add_action(RelayRLAction::minimal(i as f32, false));
}

// Full codec configuration. Requires "codec-full" (compression + encryption +
// integrity + metadata + quantization).
let key = generate_key();
let config = CodecConfig {
    compression: Some(CompressionScheme::Lz4),
    encryption_key: Some(key),
    verify_integrity: true,      // Enable BLAKE3 checksums
    include_metadata: true,
};

// Encode: Serialize → Compress → Encrypt → Checksum
let encoded = trajectory.encode(&config)?;

// Integrity is automatically verified during decode.
// `RelayRLTrajectory::decode` returns `(Self, usize)`.
let (decoded, _consumed) = RelayRLTrajectory::decode(&encoded, &config)?;

println!("Encoded {} actions", decoded.len());
println!("Total reward: {}", decoded.total_reward());

5. Chunking for Large Data

use relayrl_types::prelude::action::CodecConfig;
use relayrl_types::prelude::trajectory::RelayRLTrajectory;

let mut trajectory = RelayRLTrajectory::new(10000);
// ... add many actions ...

// `encode_chunked`/`decode_chunked` require the "metadata" and "integrity" features.
let config = CodecConfig::default();
let chunk_size = 1024 * 1024; // 1MB chunks

// Encode and split into chunks for network transmission
let chunks = trajectory.encode_chunked(&config, chunk_size)?;
println!("Split into {} chunks", chunks.len());

// ... transmit chunks over network ...

// Reassemble on the receiving end
let decoded = RelayRLTrajectory::decode_chunked(&chunks, &config)?;

6. Metadata Tracking

use relayrl_types::prelude::trajectory::RelayRLTrajectory;
use uuid::Uuid;

// Create trajectory with full metadata
let trajectory = RelayRLTrajectory::with_metadata(
    1000,                          // max_length
    Some(Uuid::new_v4()),         // agent_id
    None,                          // env_id
    None,                          // env_label
    Some(42),                      // episode number
    Some(1000),                    // training_step
);

// Check age
println!("Trajectory age: {}s", trajectory.age_seconds());

// Access metadata
if let Some(agent_id) = trajectory.get_agent_id() {
    println!("Agent: {}", agent_id);
}

Codec Pipeline

The encoding pipeline processes data in this order:

┌─────────────────┐
│  RelayRLAction  │
│ RelayRLTraject. │
└────────┬────────┘
         │
         ▼
   ┌──────────┐
   │ Bincode  │  Serialize to bytes
   └────┬─────┘
        │
        ▼
   ┌──────────┐
   │ Compress │  LZ4 or Zstd (optional)
   └────┬─────┘
        │
        ▼
   ┌──────────┐
   │ Encrypt  │  ChaCha20-Poly1305 (optional)
   └────┬─────┘
        │
        ▼
   ┌──────────┐
   │ Checksum │  BLAKE3 integrity (optional)
   └────┬─────┘
        │
        ▼
   ┌──────────┐
   │  Output  │  Final encoded bytes
   └──────────┘

Decoding reverses this pipeline with automatic verification.

ONNX Model Ingestion

ModelModule::load_from_path / ModelModule::from_onnx_bytes accept any ONNX graph that exposes exactly one tensor input and one tensor output — the graph is not required to name its input "input" or its output "output". At load time, relayrl_types:

Commits the graph with ONNX Runtime and introspects the real input/output names, element types, and shapes directly from the loaded session (Outlets), instead of assuming fixed names.
Validates the caller-supplied ModelMetadata against that discovered signature:
- input_dtype/output_dtype must map to the graph's actual ONNX element type exactly.
- input_shape/output_shape must have the same rank as the graph's declared shape.
- Every fixed graph dimension (i.e. not a dynamic/symbolic axis) must equal the corresponding metadata dimension; dynamic axes (e.g. a symbolic batch dimension) accept any concrete size at runtime.
At inference time, binds the input tensor and reads the output tensor by their discovered names, and returns the model's actual ONNX Runtime output shape (important for graphs with a dynamic batch dimension), rather than a shape reconstructed from metadata.

Graphs with more than one input/output, or with non-tensor I/O (sequences, maps, optionals), are rejected with a specific ModelError at construction time rather than failing later inside Session::run.

metadata.json itself is unchanged and remains required — RelayRL still needs it for shapes/dtypes/device defaults and as the cross-process wire contract used by the transport layer (coordinator bundles, NATS ModelFilesBundle, and on-disk directories); it is now validated against the graph instead of being trusted blindly. See crates/relayrl_types/tests/onnx_ingestion.rs for worked examples, including non-standard I/O names, dynamic batch dimensions, bool tensors, and each rejection case.

Performance Tips

LZ4: Best for real-time inference (3-4 GB/s decompression)
Zstd: Best compression ratio for training data (5-10x reduction)
Chunking: Use for trajectories > 10MB for network transmission
Integrity: Minimal overhead (~50ns per MB with BLAKE3)
Encryption: ~1 GB/s with ChaCha20-Poly1305

Examples

See the tests/ directory for more examples:

Basic action/trajectory usage
Compression benchmarks
Encryption examples
Chunked network transmission

License

Apache-2.0

relayrl_types 0.9.1