Chore/diffusion utils fix

[Kim-JeongHan]

Summary by CodeRabbit

• New Features

  • Added diffusion-based trajectory planning with guided sampling and value model integration
  • Added one-shot diffusion trajectory generation example with visualization support
  • Added RRT dataset generation script for creating training data
  • Added collision-free trajectory selection utilities
  • Added bounded collision checking for spatial constraints

• Documentation

  • Updated README with diffusion planning workflow, installation, and usage instructions

• Chores

  • Added comprehensive test coverage for diffusion modules and training pipeline
  • Added training CLI with checkpoint management and TensorBoard logging support
  • Updated project configuration with diffuser dependencies and build system setup

[coderabbitai]

Caution

Review failed

The pull request is closed.

📝 Walkthrough

Walkthrough

This PR introduces a complete diffusion-based trajectory planning pipeline, including configuration files, diffusion and value model implementations with temporal U-Net architectures, training infrastructure with checkpointing and dataset handling, sampling and inference utilities with guided policy support, checkpoint management with templating, and command-line training scripts, along with comprehensive examples and tests.

Changes

Cohort / File(s)	Summary
Configuration Files .gitignore, config/dataset_generation.yaml, config/diffusion_3d_training.yaml, config/diffusion_trajectory_one_shot.yaml, config/value_3d_training.yaml, config/__init__.py	Adds diffusion-related ignore patterns and introduces YAML configs for dataset generation, diffusion training, value training, and one-shot trajectory execution with parameters for hyperparameters, paths, and sampling settings.
Core Diffusion Models planning/diffusion/core.py, planning/diffusion/model/diffusion.py, planning/diffusion/model/policy.py, planning/diffusion/model/utils.py, planning/diffusion/model/__init__.py	Implements temporal U-Net-based DiffusionModel and ValueModel with sinusoidal embeddings, residual blocks, and encoder-decoder architectures; includes state normalization utilities (PlannerStateNormalizer, DiffusionDataset, DiffusionExperiment).
Training Infrastructure planning/diffusion/training/noise.py, planning/diffusion/training/diffusion_trainer.py, planning/diffusion/training/value_trainer.py, planning/diffusion/training/dataset.py, planning/diffusion/training/checkpoint.py, planning/diffusion/training/config.py, planning/diffusion/training/trainer.py, planning/diffusion/training/__init__.py	Establishes diffusion noise schedules (linear/cosine), epoch-level trainers for diffusion and value models, dataset loading with normalization, checkpoint persistence with EMA support, and a unified DiffusionTrainingPipeline with logging, early stopping, and multi-stage orchestration.
Sampling and Inference planning/diffusion/sampling.py, planning/diffusion/inference.py, planning/diffusion/config/__init__.py, planning/diffusion/__init__.py	Implements DiffusionSamplingEngine with inpainting-based conditioning, GuidedPolicy with optional value guidance, gradient-based guidance mechanisms, and structured configs (DiffusionConfig, DiffusionTrainingPipelineConfig); exposes unified public API surface.
Utilities and Checkpoint Management planning/diffusion/utils.py	Provides checkpoint discovery, templating context resolution with fallbacks, model reconstruction from payloads, DiffusionArtifactLoader for loading experiments, CheckpointCatalog for checkpoint enumeration, and compatibility checks between diffusion and value models.
Planning Integration planning/collision/collision_checker.py, planning/collision/__init__.py, planning/constraint/trajectory.py, planning/constraint/__init__.py, planning/sampling/diffusion_guided_sampler.py, planning/sampling/__init__.py, planning/sampling/rrt.py	Adds BoundedCollisionChecker for coordinate bounds enforcement; introduces trajectory filtering via select_collision_free_trajectory; implements DiffusionGuidedSampler for diffusion-based planner integration; refactors RRT sampler configuration to use unified sampler_kwargs dictionary.
Training and Dataset Scripts scripts/diffusion_train.py, scripts/train_diffusion.py, scripts/train_value.py, scripts/generate_rrt_dataset.py	Provides CLI entrypoints for diffusion-only and value-only training, plus RRT-based trajectory dataset generation with configurable bounds, collision checking, and NPZ output.
Example and Testing examples/diffusion_trajectory_one_shot_example.py, tests/test_diffuser_*.py, tests/test_diffusion_*.py, tests/test_rrt.py, tests/test_sampling_diffusion_guided_sampler.py	Comprehensive example showing one-shot diffusion trajectory generation with Pydantic config loading, environment setup, model loading, policy execution, trajectory filtering, and Viser visualization; unit and smoke tests for sampling, training, utilities, and integration.
Documentation and Build README.md, pyproject.toml	Expands README with Diffuser section including workflow, installation, and examples; adds build system configuration and optional diffuser dependencies (torch, pyyaml, tensorboard).

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant Config as Configuration Loader
    participant Env as Environment Setup
    participant Models as Model Loader
    participant Policy as GuidedPolicy
    participant Engine as DiffusionSamplingEngine
    participant Validator as CollisionValidator
    participant Viz as Visualizer

    User->>Config: Load YAML config
    Config->>Env: Provide environment params
    Env->>Env: Create map & obstacles
    
    User->>Models: Load checkpoints
    Models->>Models: Reconstruct DiffusionModel & ValueModel
    Models->>Policy: Initialize with models & guide
    
    User->>Policy: Request trajectory samples
    Policy->>Engine: Call sampling with conditions
    Engine->>Engine: Initialize from noise schedule
    
    loop Diffusion Steps
        Engine->>Models: Predict noise via DiffusionModel
        Engine->>Models: Compute value gradient (optional)
        Engine->>Engine: Denoise & apply guidance
    end
    
    Engine->>Policy: Return trajectory batch
    Policy->>Validator: Check collision-free segments
    Validator->>Validator: Verify start/goal tolerance
    
    alt Valid Trajectory Found
        Policy->>Viz: Return selected trajectory
        Viz->>User: Display & save visualization
    else No Valid Trajectory
        Policy->>User: Return None
    end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

Possibly related PRs

• [FEAT] Add RRT* algorithm implementation and example with mixed obsta… #4 — Modifies RRTConfig, RRTStarConfig, and sampler construction patterns, directly related to the sampler_kwargs refactoring in this PR's planning/sampling/rrt.py changes.

✨ Finishing Touches

• 📝 Generate docstrings (stacked PR)
• 📝 Generate docstrings (commit on current branch)

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch chore/diffusion-utils-fix

Tip

Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}