Software Architecture

The NRDK software architecture is designed around modular components which can be assembled by dependency injection via hydra.

In addition to the core framework and specification, the nrdk library hosts common, stable implementations of each component; experimental "research" components should be implemented in separate repositories, and only merged once stable and ready to publish.

NRDK ArchitectureGRT Reference Implementation

Core Framework

The NRDK framework is based on Pytorch Lightning, and centered around a core NRDKLightningModule.

Tip

The core framework can be completely bypassed if you wish; simply use the data loading, preprocessing, and training objectives without using the high level control flow provided in nrdk.framework.

• A model is a nn.Module which takes each batch as an input, and returns the model's predictions.

  Info

  We provide a TokenizerEncoderDecoder module, which wraps a tokenizer, encoder, and one or more decoders into a standardized model container.

• The model may or may not have GPU-side preprocessing which is expected to be performed on each batch.

• The model is trained to fit an Objective, which takes the batch and the predictions as input and returns the loss.

Data Loading

The NRDK leans heavily on the core abstract dataloader specifications, which describe a set of modular and composable data loading and preprocessing component interfaces:

Load: Sensor → Trace → Dataset
Process: Pipeline := .sample:Transform → Collate → .batch:Transform

• When using data collected by red-rover and stored in the roverd format, read data using the red-rover/roverd library, which provides abstract-dataloader compliant Dataset, Trace, and Sensor implementations.

• To preprocess this data during training, we also supply the nrdk.roverd submodule, which implements data preprocessing for each red-rover data modality; these are intended to be combined into a Transform using abstract_dataloader.ext.graph.Transform.

  Tip

  While nominally designed for red-rover data, the nrdk.roverd transforms are implemented using protocol types which allow any compatible types to be substituted as inputs.

• Data augmentations can also be specified via the abstract_dataloader.ext.augment data augmentation specifications.

Training Objectives

Training objectives for the NRDK use the abstract_dataloader.ext.objective specification. From the abstract dataloader documentation:

• An Objective is a callable which returns a (batched) scalar loss and a dictionary of metrics.
• Objectives can be combined into a higher-order objective, MultiObjective, which combines their losses and aggregates their metrics; specify these objectives using a MultiObjectiveSpec.

We currently provide the following objectives:

Objective	Ground Truth Type	Model Output Type
Occupancy2D	Occupancy2DData	Float[Tensor, 'batch t azimuth range']
Occupancy3D	Occupancy3DData	Float[Tensor, 'batch t elevation azimuth range']
Semseg	SemsegData	Float[Tensor, 'batch t h w cls']
Velocity	VelocityData	Float[Tensor, 'batch t 4']

Tip

In addition to implementing the abstract specification, each provided objective includes a type specification for their expected model predictions and ground truth data. Implementations using these objectives out-of-the-box only need to provide data which fits these type interfaces.

Model & Module Zoo

We include a number of reusable implementations of common (and stable) model modules and architectures with the nrdk; note that active research code should go in separate repositories, and only merged here once stable — and ready to publicly release!

• nrdk.models

  ---

  reusable, stable model architectures.

• nrdk.modules

  ---

  modules beyond of the standard library.

Other Modules

The nrdk includes a number of other submodules intended as reusable libraries which are not associated with a well-defined abstract specification.

• nrdk.config

  ---

  configuration utilities

• nrdk.metrics

  ---

  training and evaluation metrics

• nrdk.visualization

  ---

  visualizations utilities (e.g., for in-training visualizations)