CustomTerrainPhysicsComponent.h

Go to the documentation of this file.

// Copyright (c) 2026 Computer Vision Center (CVC) at the Universitat Autonoma
// de Barcelona (UAB).
//
// This work is licensed under the terms of the MIT license.
// For a copy, see <https://opensource.org/licenses/MIT>.
 
#pragma once
 
#include "Carla/Math/DVector.h"
#include "Carla/Vehicle/CarlaWheeledVehicle.h"
#include "Carla/MapGen/LargeMapManager.h"
#ifdef WITH_PYTORCH
THIRD_PARTY_INCLUDES_START
#include <carla/pytorch/pytorch.h>
THIRD_PARTY_INCLUDES_END
#endif
 
#include <util/ue-header-guard-begin.h>
#include "Components/ActorComponent.h"
#include "Components/PrimitiveComponent.h"
#include "Engine/StaticMeshActor.h"
#include "Engine/TextureRenderTarget2D.h"
#include "Engine/DataAsset.h"
#include "Engine/StreamableManager.h"
#include "Async/Future.h"
#include "Misc/ScopeLock.h"
#include <util/ue-header-guard-end.h>
 
#include <unordered_map>
#include <vector>
#include <string>
 
#include "CustomTerrainPhysicsComponent.generated.h"
 
 
UENUM(BlueprintType)
enum EDefResolutionType
{
  E256M = 0   UMETA(DisplayName = "256M"),
  E512M = 1   UMETA(DisplayName = "512M"),
  E1K = 2     UMETA(DisplayName = "1K"),
  E2K = 3     UMETA(DisplayName = "2K"),
};
 
UCLASS(BlueprintType)
class UHeightMapDataAsset : public UPrimaryDataAsset
{
  GENERATED_BODY()
public:
 
  UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = HeightMapDataAsset)
  int SizeX = 0;
  UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = HeightMapDataAsset)
  int SizeY = 0;
  UPROPERTY(BlueprintReadWrite, Category = HeightMapDataAsset)
  TArray<float> HeightValues;
};
 
struct FParticle
{
  FDVector Position; // position in m
  FVector Velocity;
  float Radius = 0.02f;
};
 
struct FHeightMapData
{
  void InitializeHeightmap(
      UHeightMapDataAsset* DataAsset, FDVector Size, FDVector Origin,
      FDVector Tile0, float ScaleZ);
  float GetHeight(FDVector Position) const; // get height at a given global 2d position
  void Clear();
// private:
  FDVector WorldSize;
  FDVector Offset;
  uint32_t Size_X;
  uint32_t Size_Y;
  float MinHeight = 0.0f;
  float MaxHeight = 10.0f;
  float Scale_Z = 1;
  FDVector Tile0Position;
  std::vector<float> Pixels;
};
 
struct FDenseTile
{
  FDenseTile();
  ~FDenseTile();
  FDenseTile(const FDenseTile& Origin);
  FDenseTile(FDenseTile&& Origin);
  FDenseTile& operator=(FDenseTile&& Origin);
 
  void InitializeTile(uint32_t TextureSize, float AffectedRadius, float ParticleSize, float Depth, 
      FDVector TileOrigin, FDVector TileEnd, const FString& SavePath, const FHeightMapData &HeightMap);
  std::vector<FParticle*> GetParticlesInRadius(FDVector Position, float Radius);
  void GetParticlesInRadius(FDVector Position, float Radius, std::vector<FParticle*> &ParticlesInRadius);
  void GetParticlesInBox(const FOrientedBox& OBox, std::vector<FParticle*> &ParticlesInRadius);
  void GetAllParticles(std::vector<FParticle*> &ParticlesInRadius);
  void InitializeDataStructure();
 
  void UpdateLocalHeightmap();
  std::vector<FParticle> Particles;
  std::vector<float> ParticlesHeightMap;
  std::vector<std::multiset<float,std::greater<float>>> ParticlesZOrdered;
  bool bParticlesZOrderedInitialized = false;
  FDVector TilePosition;
  FString SavePath;
  bool bHeightmapNeedToUpdate = false;
  uint32_t PartialHeightMapSize = 0;
  uint32_t TileSize = 0;
};
 
class FSparseHighDetailMap
{
public:
  friend struct FTilesWorker;
 
  FSparseHighDetailMap(float ParticleDiameter = 0.02f, float Depth = 0.4f)
    : ParticleSize(ParticleDiameter), TerrainDepth(Depth) {};
 
  void Init( uint32 NewTextureSize, float NewAffectedRadius, float ParticleDiameter, 
    float Depth, float NewFloorHeight )
  {
    ParticleSize = ParticleDiameter;
    TerrainDepth = Depth;
    TextureSize = NewTextureSize;
    AffectedRadius = NewAffectedRadius;
    FloorHeight = NewFloorHeight;
    UE_LOG(LogCarla, Warning, 
        TEXT("ParticleSize %f"), ParticleSize);
  }
 
  inline float GetTileSize() const {
    return TileSize;
  }
 
  std::vector<FParticle*> GetParticlesInRadius(FDVector Position, float Radius);
  std::vector<FParticle*> GetParticlesInTileRadius(FDVector Position, float Radius);
  std::vector<FParticle*> GetParticlesInBox(const FOrientedBox& OBox);
  std::vector<uint64_t> GetIntersectingTiles(const FOrientedBox& OBox);
  std::vector<uint64_t> GetLoadedTilesInRange(FDVector Position, float Radius);
 
 
  FDenseTile& GetTile(uint32_t Tile_X, uint32_t Tile_Y);
  FDenseTile& GetTile(FDVector Position);
  FDenseTile& GetTile(uint64_t TileId);
 
  FDenseTile& InitializeRegion(uint32_t Tile_X, uint32_t Tile_Y);
  FDenseTile& InitializeRegion(uint64_t TileId);
  FDenseTile& InitializeRegionInCache(uint64_t TileId);
 
  uint64_t GetTileId(uint32_t Tile_X, uint32_t Tile_Y);
  uint64_t GetTileId(uint64_t TileId);
  uint64_t GetTileId(FDVector Position);
  FIntVector GetVectorTileId(FDVector Position);
  FIntVector GetVectorTileId(uint64_t TileId);
  FDVector GetTilePosition(uint64_t TileId);
  FDVector GetTilePosition(uint32_t Tile_X, uint32_t Tile_Y);
 
  float GetHeight(FDVector Position) {
    return Heightmap.GetHeight(Position);
  }
 
  void InitializeMap(UHeightMapDataAsset* DataAsset,
      FDVector Origin, FDVector MapSize, float Size, float ScaleZ);
 
  void UpdateHeightMap(UHeightMapDataAsset* DataAsset,
      FDVector Origin, FDVector MapSize, float Size, float ScaleZ);
 
  void UpdateMaps(FDVector Position, float RadiusX, float RadiusY, float CacheRadiusX, float CacheRadiusY);
 
  void Update(FVector Position, float RadiusX, float RadiusY);
 
  void SaveMap();
 
  void Clear();
 
  void LockMutex()
  {
    Lock_Map.Lock();
  }
 
  void UnLockMutex()
  {
    Lock_Map.Unlock();
  }
 
  std::vector<uint64_t> GetTileIdInMap()
  {
    std::vector<uint64_t> Result;
    for (auto& Iter : Map)
    {
      Result.emplace_back(Iter.first);
    }
    return Result;
  }
  std::vector<uint64_t> GetTileIdInCache()
  {
    std::vector<uint64_t> Result;
    for (auto& Iter : CacheMap)
    {
      Result.emplace_back(Iter.first);
    }
    return Result;
  }
 
  std::unordered_map<uint64_t, FDenseTile> Map;
  std::unordered_map<uint64_t, FDenseTile> CacheMap;
  FString SavePath;
  FCriticalSection Lock_Particles;
private:
  std::unordered_map<uint64_t, FDenseTile> TilesToWrite;
  FDVector Tile0Position;
  FDVector Extension;
  float TileSize = 1.f; // 1m per tile
  FHeightMapData Heightmap;
  float ParticleSize = 0.02f;
  float TerrainDepth = 0.4f;
  float FloorHeight = 0.0f;
  uint32 TextureSize = 0;
  float AffectedRadius = 0.0f;
  FVector PositionToUpdate;
  FCriticalSection Lock_Map; // UE4 Mutex
  FCriticalSection Lock_CacheMap; // UE4 Mutex
  FCriticalSection Lock_GetTile;
  FCriticalSection Lock_Position; // UE4 Mutex
 
};
 
USTRUCT(BlueprintType)
struct FForceAtLocation
{
  GENERATED_BODY()
 
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  FVector Force = FVector::ZeroVector;
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  FVector Location = FVector::ZeroVector;
};
 
UCLASS( ClassGroup=(Custom), meta=(BlueprintSpawnableComponent) )
class CARLA_API UCustomTerrainPhysicsComponent : public UActorComponent
{
  GENERATED_BODY()
  friend struct FTilesWorker;
public:
 
  UCustomTerrainPhysicsComponent();
  
  virtual void BeginPlay() override;
  virtual void EndPlay(const EEndPlayReason::Type EndPlayReason) override;
 
  virtual void TickComponent(float DeltaTime,
      ELevelTick TickType, FActorComponentTickFunction* ThisTickFunction);
 
  // UFUNCTION(BlueprintCallable)
  // TArray<FHitResult> SampleTerrainRayCast(const TArray<FVector> &Locations);
 
  UFUNCTION(BlueprintCallable)
  void AddForces(const TArray<FForceAtLocation> &Forces);
 
  UFUNCTION(BlueprintCallable)
  static void BuildLandscapeHeightMapDataAasset(ALandscapeProxy* Landscape, 
      int Resolution, FVector MapSize, FString AssetPath, FString AssetName);
 
  UFUNCTION(BlueprintCallable)
  float GetHeightAtLocation(ALandscapeProxy * Landscape, FVector Location);
 
  UFUNCTION(BlueprintCallable)
  TArray<FVector> GetParticlesInRadius(FVector Position, float Radius);
 
  UFUNCTION(BlueprintCallable)
  TArray<FVector> GetParticlesInTileRadius(FVector Position, float Radius);
 
  UFUNCTION(BlueprintCallable)
  FVector GetTileCenter(FVector Position);
 
  UFUNCTION(BlueprintCallable, Category="Tiles")
  void UpdateMaps(FVector Position, float RadiusX, float RadiusY, float CacheRadiusX, float CacheRadiusY);
 
  UFUNCTION(BlueprintCallable, Category="Texture")
  void InitTexture();
 
  UFUNCTION(BlueprintCallable, Category="Texture")
  void UpdateTexture();
 
  UFUNCTION(BlueprintCallable, Category="Texture")
  void UpdateLoadedTextureDataRegions();
  
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  UHeightMapDataAsset* DataAsset;
 
  UFUNCTION(BlueprintCallable, Category="Texture")
  void UpdateLargeTexture();
  
  UFUNCTION(BlueprintCallable, Category="Texture")
  void UpdateLargeTextureData();
 
  UPROPERTY(EditAnywhere, BlueprintReadWrite, Category="MaterialParameters")
  UTexture2D* TextureToUpdate;
  UPROPERTY(EditAnywhere, BlueprintReadWrite, Category="MaterialParameters")
  float MinDisplacement = -100;
  UPROPERTY(EditAnywhere, BlueprintReadWrite, Category="MaterialParameters")
  float MaxDisplacement = 100;
 
  UPROPERTY(EditAnywhere, BlueprintReadWrite, Category="MaterialParameters")
  UTexture2D* LargeTextureToUpdate;
 
  UPROPERTY(EditAnywhere, BlueprintReadWrite, Category="MaterialParameters")
  UMaterialParameterCollection* MPC;
 
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  FString NeuralModelFile = "";
 
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  FVector NextPositionToUpdate = FVector(0,0,0);
  
  FVector LastUpdatedPosition;
  FVector CachePosition;
 
  FString SavePath;
 
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  float ForceMulFactor = 1.0;
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  float ParticleForceMulFactor = 1.0;
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  int SoilType = 0;
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  bool bUseSoilType = false;
  UPROPERTY(EditAnywhere)
  bool NNVerbose = false;
  UPROPERTY(EditAnywhere, BlueprintReadWrite)
  bool bUseLocalFrame = false;
private:
 
  void RunNNPhysicsSimulation(
      ACarlaWheeledVehicle *Vehicle, float DeltaTime);
  // TArray<FParticle*> GetParticlesInRange(...);
  void SetUpParticleArrays(std::vector<FParticle*>& ParticlesIn, 
      TArray<float>& ParticlePosOut, 
      TArray<float>& ParticleVelOut,
      const FTransform &WheelTransform);
  void SetUpWheelArrays(ACarlaWheeledVehicle *Vehicle, int WheelIdx,
      TArray<float>& WheelPos, 
      TArray<float>& WheelOrientation, 
      TArray<float>& WheelLinearVelocity, 
      TArray<float>& WheelAngularVelocity);
  void UpdateParticles(
      std::vector<FParticle*> Particles, std::vector<float> Forces,
      float DeltaTime, const FTransform& WheelTransform);
  void ApplyForcesToVehicle(
      ACarlaWheeledVehicle *Vehicle,
      FVector ForceWheel0, FVector ForceWheel1, FVector ForceWheel2, FVector ForceWheel3,
      FVector TorqueWheel0, FVector TorqueWheel1, FVector TorqueWheel2, FVector TorqueWheel3);
  void ApplyMeanAccelerationToVehicle(
      ACarlaWheeledVehicle *Vehicle,
      FVector ForceWheel0, FVector ForceWheel1, FVector ForceWheel2, FVector ForceWheel3);
  void ApplyAccelerationToVehicle(
      ACarlaWheeledVehicle *Vehicle,
      FVector ForceWheel0, FVector ForceWheel1, FVector ForceWheel2, FVector ForceWheel3);
      
  void ApplyForces();
  void LimitParticlesPerWheel(std::vector<FParticle*> &Particles);
  void DrawParticles(UWorld* World, std::vector<FParticle*>& Particles, 
      FLinearColor Color = FLinearColor(1.f, 0.f, 0.f));
  void DrawParticlesArray(UWorld* World, TArray<float>& ParticlesArray, 
      FLinearColor Color = FLinearColor(1.f, 0.f, 0.f));
  void DrawOrientedBox(UWorld* World, const TArray<FOrientedBox>& Boxes);
  void DrawTiles(UWorld* World, const std::vector<uint64_t>& TilesIds, float Height = 0,
    FLinearColor Color = FLinearColor(0.0,1.0,0.0,1.0));
  void GenerateBenchmarkParticles(std::vector<FParticle>& BenchParticles, 
      std::vector<FParticle*> &ParticlesWheel0, std::vector<FParticle*> &ParticlesWheel1,
      std::vector<FParticle*> &ParticlesWheel2, std::vector<FParticle*> &ParticlesWheel3,
      FOrientedBox &BboxWheel0, FOrientedBox &BboxWheel1, 
      FOrientedBox &BboxWheel2, FOrientedBox &BboxWheel3);
 
  void UpdateParticlesDebug(std::vector<FParticle*> Particles);
  
  void OnLevelAddedToWorld(ULevel* InLevel, UWorld* InWorld);
  
  void UpdateTilesHeightMaps( const std::vector<FParticle*>& Particles);
  void RemoveParticlesFromOrderedContainer(const std::vector<FParticle*>& Particles);
  void AddParticlesToOrderedContainer(const std::vector<FParticle*>& Particles);
  void FlagTilesToRedoOrderedContainer(const std::vector<FParticle*>& Particles);
  void UpdateTilesHeightMapsInRadius(FDVector Position, uint32 Rad );
  
  void AddForceToSingleWheel(USkeletalMeshComponent* SkeletalMeshComponent,
    FVector WheelPosition, FVector WheelNormalForce);
  
  UPROPERTY(EditAnywhere)
  TArray<FForceAtLocation> ForcesToApply;
  UPROPERTY(EditAnywhere)
  TObjectPtr<UPrimitiveComponent> RootComponent;
  UPROPERTY(EditAnywhere)
  float RayCastRange = 10.0f;
 
public:
  UPROPERTY(EditAnywhere)
  FVector WorldSize = FVector(200000,200000,0);
  UPROPERTY(EditAnywhere)
  bool DrawDebugInfo = true;
  UPROPERTY(EditAnywhere)
  bool bUpdateParticles = false;
  // Radius of the data loaded in memory
  UPROPERTY(EditAnywhere, Category="Tiles")
  FVector TileRadius = FVector( 100, 100, 0 );
  // Radius of the data loaded in memory
  UPROPERTY(EditAnywhere, Category="Tiles")
  FVector CacheRadius = FVector( 50, 50, 0 );
  UPROPERTY(EditAnywhere, Category="Tiles")
  bool bDrawLoadedTiles = false;
  UPROPERTY(EditAnywhere, Category="Tiles")
  int32 TileSize = 1;
  UPROPERTY(EditAnywhere, Category="Tiles")
  bool bRemoveLandscapeColliders = false;
private:
  // TimeToTriggerCacheReload In seconds
  UPROPERTY(EditAnywhere, Category="Tiles")
  float TimeToTriggerCacheReload = 20.0f;
  // TimeToTriggerLoadTiles in MS
  UPROPERTY(EditAnywhere, Category="Tiles")
  float TimeToTriggerLoadTiles = 1.0f;
  UPROPERTY(EditAnywhere, Category="Tiles")
  float TimeToTriggerUnLoadTiles = 5.0f;
  // Radius of the data collected by the texture in METERS
  UPROPERTY(EditAnywhere, Category="MaterialParameters")
  float TextureRadius = 4.0f;
  // Radius of the data collected by the texture in METERS
  UPROPERTY(EditAnywhere, Category="MaterialParameters")
  float LargeTextureRadius = 50.0f;
  // Scalar Factor of deformation effect applied in the landscape
  UPROPERTY(EditAnywhere, Category="MaterialParameters")
  float EffectMultiplayer = 10.0f;
 
  UPROPERTY(EditAnywhere, Category="MaterialParameters")
  TEnumAsByte<EDefResolutionType> ChosenRes = EDefResolutionType::E1K;
  UPROPERTY(EditAnywhere, Category="MaterialParameters")
  TMap<TEnumAsByte<EDefResolutionType>, UTexture2D*> TexturesRes;
 
  bool bVisualization = false;
 
  UPROPERTY(EditAnywhere, Category="DeformationMesh")
  bool bUseDeformationPlane = false;
  UPROPERTY(EditAnywhere, Category="DeformationMesh")
  TObjectPtr<UStaticMesh> DeformationPlaneMesh = nullptr;
  UPROPERTY(EditAnywhere, Category="DeformationMesh")
  TObjectPtr<UMaterialInstance> DeformationPlaneMaterial = nullptr;
  UPROPERTY(VisibleAnywhere, Category="DeformationMesh")
  TObjectPtr<AStaticMeshActor> DeformationPlaneActor = nullptr;
 
  UPROPERTY()
  TObjectPtr<UMaterialParameterCollectionInstance> MPCInstance;
 
  UPROPERTY(EditAnywhere, Category="Forces")
  float NormalForceIntensity = 100;
  
  UPROPERTY(EditAnywhere)
  float SearchRadius = 100;
  UPROPERTY(EditAnywhere)
  float ParticleDiameter = 2;
  UPROPERTY(EditAnywhere)
  float TerrainDepth = 40;
  UPROPERTY(EditAnywhere)
  TObjectPtr<AActor> FloorActor = nullptr;
  UPROPERTY(EditAnywhere)
  bool bUseDynamicModel = false;
  UPROPERTY(EditAnywhere)
  bool bUseCUDAModel = false;
 
  UPROPERTY(EditAnywhere)
  float TireRadius = 33.0229f;
  UPROPERTY(EditAnywhere)
  float TireWidth = 21.21f;
  UPROPERTY(EditAnywhere)
  float BoxSearchForwardDistance = 114.39f;
  UPROPERTY(EditAnywhere)
  float BoxSearchLateralDistance = 31.815f;
  UPROPERTY(EditAnywhere)
  float BoxSearchDepthDistance = 20.f;
  UPROPERTY(EditAnywhere)
  bool bDisableVehicleGravity = false;
  UPROPERTY(EditAnywhere)
  float MaxForceMagnitude = 1000000.f;
  UPROPERTY(EditAnywhere)
  float FloorHeight = 0.f;
  UPROPERTY(EditAnywhere)
  bool bUseImpulse = false;
  UPROPERTY(EditAnywhere)
  bool bUseMeanAcceleration = false;
  UPROPERTY(EditAnywhere)
  bool bShowForces = true;
  UPROPERTY(EditAnywhere)
  float MinHeight = 0;
  UPROPERTY(EditAnywhere)
  float MaxHeight = 10;
  UPROPERTY(EditAnywhere)
  FVector Tile0Origin;
  UPROPERTY(EditAnywhere)
  bool bDrawHeightMap = false;
  UPROPERTY(EditAnywhere)
  FVector DrawStart = FVector(0);
  UPROPERTY(EditAnywhere)
  FVector DrawEnd = FVector(1000, 1000, 0);
  UPROPERTY(EditAnywhere)
  FVector DrawInterval = FVector(100,100,0);
  UPROPERTY(EditAnywhere)
  int CUDADevice = 0;
  UPROPERTY(EditAnywhere)
  FVector HeightMapScaleFactor = FVector(1, 1, 1);
  UPROPERTY(EditAnywhere)
  FVector HeightMapOffset = FVector(0, 0, 0);
  UPROPERTY(EditAnywhere)
  bool bBenchMark = false;
  UPROPERTY(EditAnywhere)
  int MaxParticlesPerWheel = 6000;
 
  UPROPERTY(EditAnywhere)
  FVector Radius = FVector(10,10,10);
  
  UPROPERTY(VisibleAnywhere)
  FIntVector CurrentLargeMapTileId = FIntVector(-1,-1,0);
  UPROPERTY(VisibleAnywhere)
  TObjectPtr<ALargeMapManager> LargeMapManager = nullptr;
 
  // Async heightmap streaming state.
  // ActiveHeightMap keeps the most recently applied data asset rooted for GC
  // while the streaming handle is the primary owner during the in-flight load.
  UPROPERTY(Transient, DuplicateTransient)
  TObjectPtr<UHeightMapDataAsset> ActiveHeightMap = nullptr;
  TSharedPtr<FStreamableHandle> PendingHeightMapHandle;
  FIntVector PendingHeightMapTileId = FIntVector(-1, -1, 0);
 
  void OnHeightMapLoaded();
 
  UPROPERTY()
  TArray<TObjectPtr<ACarlaWheeledVehicle>> Vehicles;
  FSparseHighDetailMap SparseMap;
  TArray<uint8> Data;
  TArray<uint8> LargeData;
  #ifdef WITH_PYTORCH
  carla::learning::NeuralModel TerramechanicsModel;
  #endif
 
  TFuture<bool> IterationCompleted;
 
  class FRunnableThread* Thread;
  struct FTilesWorker* TilesWorker;
};
 
struct FTilesWorker : public FRunnable
{
    FTilesWorker(class UCustomTerrainPhysicsComponent* TerrainComp, FVector NewPosition, float NewRadiusX, float NewRadiusY );
 
    virtual ~FTilesWorker() override;
 
    virtual uint32 Run() override; 
 
  class UCustomTerrainPhysicsComponent* CustomTerrainComp;
  FVector Position;
 
  double RadiusX; 
  double RadiusY;
 
  volatile bool bShouldContinue = true;
};