This document provides guidance for AI agents working with the Neural Radar Development Kit (NRDK), a Python framework for developing, training, and evaluating machine learning models on radar spectrum and multimodal sensor data.
The NRDK is built around modular components that can be assembled via dependency injection using Hydra configuration. The framework follows a strict separation of concerns:
- Core Framework: PyTorch Lightning-based training pipeline
- Data Loading: Abstract dataloader specifications for modular data processing
- Training Objectives: Standardized loss functions and metrics
- Model Zoo: Reusable, stable model architectures and modules
Components should implement well-defined interfaces and be composable. Use the abstract dataloader specifications for data processing and the objective interface for training losses.
All code should be fully typed using
jaxtyping
and
beartype
. Type annotations are both statically checked and enforced at runtime.
Use Hydra configuration files for all experiments. Configuration should be hierarchical and composable.
Only merge stable, production-ready code into the main repository. Research code should live in separate repositories until proven stable.
When adding new data preprocessing:
- Implement Abstract Interfaces: Use
abstract_dataloader.spec
interfaces (Sensor
, Trace
, Dataset
, Transform
, Collate
)
- Protocol Types: Leverage protocol types for compatibility with existing transforms
- Configuration: Create corresponding Hydra config files in the appropriate config group
- Documentation: Include type specifications for inputs and outputs
When adding new models:
- Extend Standard Modules: Build on
torch.nn.Module
or use TokenizerEncoderDecoder
wrapper
- Size Configuration: Reference global size parameters from the
size
config group
- Component Structure: Organize as tokenizer/encoder/decoder components when applicable
- Type Annotations: Specify exact tensor shapes using
jaxtyping
When implementing new objectives:
- Follow Objective Interface: Implement
abstract_dataloader.ext.objective.Objective
- Return Loss and Metrics: Return scalar loss and dictionary of metrics
- Type Ground Truth: Specify expected input/output types clearly
- Support Multi-Objectives: Ensure compatibility with
MultiObjective
wrapper
When extending CLI functionality:
- Use Tyro Framework: Follow existing patterns with positional and flagged arguments
- Validation: Include validation functions for configurations and data
- Export/Import: Support standard export formats for interoperability
- Help Documentation: Provide clear usage examples and parameter descriptions
NRDK uses pytest for automated testing. To run the test suite:
# Run all tests
uv run --extra dev pytest tests
# Run tests with coverage report
uv run --extra dev pytest --cov=src/nrdk tests
# Run specific test file
uv run --extra dev pytest tests/test_basic.py
# Run with verbose output
uv run --extra dev pytest -v tests
When contributing to NRDK, include tests for new functionality:
- Test Structure: Place tests in the
tests/
directory
- Naming Convention: Use
test_*.py
for test files and test_*
for test functions
- Test Categories:
- Unit tests: Test individual components in isolation
- Integration tests: Test component interactions
- Configuration tests: Validate Hydra configuration loading
Example test structure:
import pytest
import nrdk
def test_component_functionality():
"""Test basic component functionality."""
# Test implementation
assert True
class TestComponentClass:
"""Test class for comprehensive component testing."""
def test_initialization(self):
"""Test component initialization."""
# Test implementation
pass
- Type Checking: All code must pass strict type checking with
pyright
- Linting: Code must pass
ruff
linting checks
- Runtime Validation: Use
beartype
for runtime type validation
- Import Hooks: Leverage jaxtyping import hooks for automatic validation
- Test Coverage: Maintain reasonable test coverage for critical components
The project uses pre-commit hooks to ensure code quality:
# Install pre-commit hooks
uv run --extra dev pre-commit install
# Run all hooks manually
uv run --extra dev pre-commit run --all-files
Hooks include:
- ruff: Code linting and formatting
- pyright: Static type checking
- pytest: Automated test execution
Use the
nrdk validate
CLI tool to:
- Validate configuration file syntax
- Check parameter compatibility
- Verify data path existence
- Ensure objective/model alignment
- Sensor Data: Validate sensor data formats and synchronization
- Transform Pipelines: Test data preprocessing chains
- Batch Processing: Verify collation and batching behavior
- Clone/Submodule: Add NRDK as submodule to your project
- Dependencies: Install with appropriate extras (
roverd
, grt
, docs
, dev
)
- Configuration: Copy GRT reference configuration as starting point
- Virtual Environment: Use
uv
for dependency management
- Configuration First: Create Hydra configs before implementing
- Incremental Development: Build components incrementally with validation
- Logging: Use rich logging with appropriate log levels
- Checkpointing: Leverage PyTorch Lightning's automatic checkpointing
- Interface Compliance: Ensure components follow abstract specifications
- Configuration Testing: Test configuration composition and overrides
- Performance Monitoring: Use included performance monitoring callbacks
- Documentation: Update type specifications and usage examples
# Standard config pattern
_target_: module.path.ClassName
param1: value1
param2: ${size.d_model} # Reference global sizes
nested_config:
_target_: another.module.Class
- Always use complete type annotations
- Leverage jaxtyping for tensor shapes
- Use beartype for runtime validation
- Keep configs small and focused
- Use composition over large monolithic configs
- Reference global parameters where appropriate
- Provide clear error messages for configuration issues
- Validate inputs at component boundaries
- Use logging for debugging information
- Use PyTorch Lightning's optimization features
- Implement efficient data loading patterns
- Monitor memory usage in long training runs
When extending NRDK functionality, consider these related projects:
- red-rover: Multimodal radar spectrum ecosystem
- xwr: Python interface for TI mmWave radars
- abstract-dataloader: Core data loading specifications
Keep these dependencies minimal and well-defined to maintain NRDK's modularity.