You've already forked LuckyWorld
Merge pull request '[FT] Remote control + jsonl files written' (#84) from rel-data-record into main
Reviewed-on: LuckyRobots/LuckyWorldV2#84
This commit is contained in:
Binary file not shown.
Binary file not shown.
BIN
Content/Levels/kitchenLevel/BP-Sensor.uasset
Normal file
BIN
Content/Levels/kitchenLevel/BP-Sensor.uasset
Normal file
Binary file not shown.
Binary file not shown.
@ -2,31 +2,19 @@
|
||||
|
||||
|
||||
#include "LuckyDataTransferSubsystem.h"
|
||||
|
||||
#include "AutomationBlueprintFunctionLibrary.h"
|
||||
#include "ImageUtils.h"
|
||||
#include "RenderingThread.h"
|
||||
#include "RenderUtils.h"
|
||||
#include "RenderGraphUtils.h"
|
||||
#include "RHI.h"
|
||||
#include "RHICommandList.h"
|
||||
#include "ImageWriteQueue.h"
|
||||
#include "ImageWriteTask.h"
|
||||
#include "ImagePixelData.h"
|
||||
#include "JsonUtilities.h"
|
||||
#include "JsonObjectConverter.h"
|
||||
#include "ReviewComments.h"
|
||||
#include "WebSocketsModule.h"
|
||||
#include "IWebSocket.h"
|
||||
#include "Kismet/KismetStringLibrary.h"
|
||||
#include "Camera/CameraActor.h"
|
||||
#include "Camera/CameraComponent.h"
|
||||
#include "Components/SceneCaptureComponent2D.h"
|
||||
#include "Kismet/GameplayStatics.h"
|
||||
#include "Kismet/KismetMathLibrary.h"
|
||||
#include "Kismet/KismetRenderingLibrary.h"
|
||||
#include "Slate/SceneViewport.h"
|
||||
#include "Virtualization/VirtualizationTypes.h"
|
||||
|
||||
ULuckyDataTransferSubsystem::ULuckyDataTransferSubsystem()
|
||||
{
|
||||
@ -89,16 +77,16 @@ void ULuckyDataTransferSubsystem::Callback_OnConnected()
|
||||
UE_LOG(LogTemp, Warning, TEXT("WebSocket connected successfully"));
|
||||
}
|
||||
|
||||
void ULuckyDataTransferSubsystem::Callback_OnConnectionError(const FString& Error)
|
||||
void ULuckyDataTransferSubsystem::Callback_OnConnectionError(const FString& Error) const
|
||||
{
|
||||
UE_LOG(LogTemp, VeryVerbose, TEXT("Websocket connection error: %s"), *Error)
|
||||
}
|
||||
|
||||
void ULuckyDataTransferSubsystem::Callback_OnMessage(const FString& Message)
|
||||
void ULuckyDataTransferSubsystem::Callback_OnMessage(const FString& Message) const
|
||||
{
|
||||
if (!Message.IsEmpty())
|
||||
{
|
||||
CommandReady(InterpretData(Message));
|
||||
CommandReady(ParseJsonPayload(Message));
|
||||
return;
|
||||
}
|
||||
|
||||
@ -130,9 +118,9 @@ void ULuckyDataTransferSubsystem::SendMessage(const FString& Message)
|
||||
UE_LOG(LogTemp, Warning, TEXT("WebSocket outgoing message failed"));
|
||||
}
|
||||
|
||||
FPayload ULuckyDataTransferSubsystem::InterpretData(const FString& Message)
|
||||
FRemoteControlPayload ULuckyDataTransferSubsystem::ParseJsonPayload(const FString& Message)
|
||||
{
|
||||
FPayload Payload = FPayload();
|
||||
FRemoteControlPayload Payload = FRemoteControlPayload();
|
||||
|
||||
if (!Message.IsEmpty())
|
||||
{
|
||||
@ -143,9 +131,9 @@ FPayload ULuckyDataTransferSubsystem::InterpretData(const FString& Message)
|
||||
{
|
||||
for (auto& Elem : JsonObj->Values)
|
||||
{
|
||||
FCommand Command = FCommand();
|
||||
Command.Key = FString(Elem.Key);
|
||||
Command.Value = Elem.Value->AsNumber();
|
||||
FRemoteControlActuatorCommand Command = FRemoteControlActuatorCommand();
|
||||
Command.ActuatorName = FString(Elem.Key);
|
||||
Command.ActuatorValue = Elem.Value->AsNumber();
|
||||
|
||||
Payload.Commands.Add(Command);
|
||||
}
|
||||
@ -161,7 +149,7 @@ void IncomingMessage(const FString& Message)
|
||||
UE_LOG(LogTemp, Warning, TEXT("Incoming message: %s"), *outMessage);
|
||||
}
|
||||
|
||||
void ULuckyDataTransferSubsystem::CommandReady(const FPayload& Payload)
|
||||
void ULuckyDataTransferSubsystem::CommandReady(const FRemoteControlPayload& Payload) const
|
||||
{
|
||||
if (OnCommandReady.IsBound())
|
||||
{
|
||||
@ -204,7 +192,7 @@ void ULuckyDataTransferSubsystem::RegisterSensor(ALuckySensorPawnBase* Sensor)
|
||||
}
|
||||
}
|
||||
|
||||
bool ULuckyDataTransferSubsystem::WriteImageToDisk(const FString& inPath, const double Timestamp, const FString& Comment)
|
||||
bool ULuckyDataTransferSubsystem::WriteImageToDisk(const double Timestamp, const FString& InPath, const FString& Comment)
|
||||
{
|
||||
if (SessionID.IsEmpty())
|
||||
{
|
||||
@ -215,79 +203,86 @@ bool ULuckyDataTransferSubsystem::WriteImageToDisk(const FString& inPath, const
|
||||
// TODO JB: I have this line in my project because I had a discrepancy in editor vs packaged
|
||||
// TODO UKismetSystemLibrary::GetProjectSavedDirectory()
|
||||
// TODO Check which one is working in shipping build
|
||||
FString Path = inPath.IsEmpty() ? FPaths::ProjectSavedDir() : inPath; // swap this for const path
|
||||
FString Path = InPath.IsEmpty() ? FPaths::ProjectSavedDir() : InPath; // swap this for const path
|
||||
|
||||
if (!SensorPawns.IsEmpty())
|
||||
if (SensorPawns.IsEmpty())
|
||||
{
|
||||
for (const ALuckySensorPawnBase* Sensor : SensorPawns)
|
||||
UE_LOG(LogTemp, Warning, TEXT("Attempted to capture image data but no sensors are registered"));
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const ALuckySensorPawnBase* Sensor : SensorPawns)
|
||||
{
|
||||
if (IsValid(Sensor) && Sensor->SensorInfo.bActive && Sensor->GetCamera() && Sensor->GetCaptureComponent())
|
||||
{
|
||||
if (IsValid(Sensor) && Sensor->SensorInfo.bActive && Sensor->GetCamera() && Sensor->GetCaptureComponent())
|
||||
// TODO Pass the robot name in params - use RobotType and EnumToString
|
||||
FString Robot = TEXT("Robot_Name");
|
||||
|
||||
FString Episode = SessionID;
|
||||
|
||||
ENQUEUE_RENDER_COMMAND(ReadPixelsAsync)([Sensor, Path, Timestamp, Comment, Episode](FRHICommandListImmediate& RHICmdList)
|
||||
{
|
||||
FString Robot = TEXT("Robot_Name");
|
||||
FTextureResource* Resource = Sensor->RenderTarget->GetResource();
|
||||
FRHITexture* ResourceRHI = Resource->GetTexture2DRHI();
|
||||
|
||||
FString Episode = SessionID;
|
||||
TArray<FColor> OutPixels;
|
||||
|
||||
ENQUEUE_RENDER_COMMAND(ReadPixelsAsync)([Sensor, Path, Timestamp, Comment, Episode](FRHICommandListImmediate& RHICmdList)
|
||||
// Fail and exit
|
||||
if (!ensure(ResourceRHI))
|
||||
{
|
||||
FTextureResource* Resource = Sensor->RenderTarget->GetResource();
|
||||
FRHITexture* ResourceRHI = Resource->GetTexture2DRHI();
|
||||
UE_LOG(LogTemp, Warning, TEXT("ResourceRHI not found"));
|
||||
return;
|
||||
}
|
||||
|
||||
TArray<FColor> OutPixels;
|
||||
OutPixels.SetNum(Resource->GetSizeX() * Resource->GetSizeY());
|
||||
RHICmdList.ReadSurfaceData(
|
||||
ResourceRHI,
|
||||
FIntRect(0, 0, Resource->GetSizeX(), Resource->GetSizeY()),
|
||||
OutPixels,
|
||||
FReadSurfaceDataFlags()
|
||||
);
|
||||
|
||||
if (ensure(ResourceRHI))
|
||||
{
|
||||
OutPixels.SetNum(Resource->GetSizeX() * Resource->GetSizeY());
|
||||
RHICmdList.ReadSurfaceData(
|
||||
ResourceRHI,
|
||||
FIntRect(0, 0, Resource->GetSizeX(), Resource->GetSizeY()),
|
||||
OutPixels,
|
||||
FReadSurfaceDataFlags()
|
||||
);
|
||||
// TODO This logs 1, it doesn't work?
|
||||
// UE_LOG(LogTemp, Warning, TEXT("Logged pixels: %d"), OutPixels.Num())
|
||||
|
||||
UE_LOG(LogTemp, Warning, TEXT("Logged pixels: %d"), OutPixels.Num())
|
||||
FImageWriteTask* ImageTask = new FImageWriteTask();
|
||||
|
||||
FImageWriteTask* ImageTask = new FImageWriteTask();
|
||||
// TODO What's that for?
|
||||
if (Path.Right(1) == "/")
|
||||
{
|
||||
//Path = Path.Left(Path.Len() - 1);
|
||||
}
|
||||
|
||||
if (Path.Right(1) == "/")
|
||||
{
|
||||
//Path = Path.Left(Path.Len() - 1);
|
||||
}
|
||||
// TODO Replace by what?
|
||||
const FString RobotName = TEXT("Robot_Name");
|
||||
|
||||
FString Robot = TEXT("Robot_Name");
|
||||
const FString Filename = FString::Printf(
|
||||
TEXT("%s/TrainingData/%s/%s/%s/%s_%s_%s_%d"),
|
||||
*Path,
|
||||
*RobotName,
|
||||
*Episode,
|
||||
*Sensor->SensorInfo.Name,
|
||||
*RobotName,
|
||||
*Sensor->SensorInfo.Name,
|
||||
*Comment,
|
||||
FMath::Floor<int32>(Timestamp * 1000)
|
||||
);
|
||||
|
||||
const FString Filename = FString::Printf(
|
||||
TEXT("LuckRobotData/%s/%s/%s/%s/%s_%s_%s_%d"),
|
||||
*Path,
|
||||
*Robot,
|
||||
*Episode,
|
||||
*Sensor->SensorInfo.Name,
|
||||
*Robot,
|
||||
*Sensor->SensorInfo.Name,
|
||||
*Comment,
|
||||
FMath::Floor<int32>(Timestamp * 1000)
|
||||
);
|
||||
UE_LOG(LogTemp, VeryVerbose, TEXT("FileName %s"), *Filename);
|
||||
|
||||
//UE_LOG(LogTemp, Warning, TEXT("Evan requested a longer string describing the inner workings of the following string which describes in great detail the file path for the image you've just written to disk. It is: %s"), *Filename);
|
||||
ImageTask->Format = EImageFormat::JPEG;
|
||||
ImageTask->Filename = Filename;
|
||||
ImageTask->PixelData = MakeUnique<TImagePixelData<FColor>>(FIntPoint(Sensor->RenderTarget->GetSurfaceWidth(), Sensor->RenderTarget->GetSurfaceHeight()), TArray64<FColor>(OutPixels));
|
||||
ImageTask->bOverwriteFile = true;
|
||||
ImageTask->CompressionQuality = static_cast<int32>(EImageCompressionQuality::Default);
|
||||
|
||||
ImageTask->Format = EImageFormat::PNG;
|
||||
ImageTask->Filename = Filename;
|
||||
ImageTask->PixelData = MakeUnique<TImagePixelData<FColor>>(FIntPoint(Sensor->RenderTarget->GetSurfaceWidth(), Sensor->RenderTarget->GetSurfaceHeight()), TArray64<FColor>(OutPixels));
|
||||
ImageTask->bOverwriteFile = true;
|
||||
ImageTask->CompressionQuality = (int32)EImageCompressionQuality::Default;
|
||||
// Add to write queue (async)
|
||||
FModuleManager::LoadModuleChecked<IImageWriteQueueModule>("ImageWriteQueue").GetWriteQueue().Enqueue(TUniquePtr<IImageWriteTaskBase>(ImageTask));
|
||||
|
||||
//Add to write queue (async)
|
||||
FModuleManager::LoadModuleChecked<IImageWriteQueueModule>("ImageWriteQueue").GetWriteQueue().Enqueue(TUniquePtr<IImageWriteTaskBase>(ImageTask));
|
||||
|
||||
ImageTask->OnCompleted = [](bool bSuccess) {
|
||||
UE_LOG(LogTemp, Warning, TEXT("Image write completed: %s"), bSuccess ? TEXT("Success") : TEXT("Failed"));
|
||||
};
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
});
|
||||
}
|
||||
ImageTask->OnCompleted = [](bool bSuccess) {
|
||||
UE_LOG(LogTemp, VeryVerbose, TEXT("Image write completed: %s"), bSuccess ? TEXT("Success") : TEXT("Failed"));
|
||||
};
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -29,7 +29,7 @@ ALuckySensorPawnBase::ALuckySensorPawnBase()
|
||||
|
||||
RenderTarget = CreateDefaultSubobject<UTextureRenderTarget2D>(TEXT("RenderTarget"));
|
||||
RenderTarget->UpdateResourceImmediate(true);
|
||||
RenderTarget->ResizeTarget(1, 1);
|
||||
RenderTarget->ResizeTarget(640, 360);
|
||||
|
||||
|
||||
CaptureComponent->CaptureSource = SCS_FinalColorLDR;
|
||||
|
||||
@ -13,32 +13,32 @@
|
||||
*
|
||||
*/
|
||||
USTRUCT(BlueprintType)
|
||||
struct FCommand
|
||||
struct FRemoteControlActuatorCommand
|
||||
{
|
||||
GENERATED_BODY()
|
||||
|
||||
public:
|
||||
UPROPERTY(BlueprintReadOnly, Category = "Command")
|
||||
FString Key = FString();
|
||||
FString ActuatorName = FString();
|
||||
|
||||
UPROPERTY(BlueprintReadOnly, Category = "Command")
|
||||
float Value = 0.f;
|
||||
float ActuatorValue = 0.f;
|
||||
};
|
||||
|
||||
USTRUCT(BlueprintType)
|
||||
struct FPayload
|
||||
struct FRemoteControlPayload
|
||||
{
|
||||
GENERATED_BODY()
|
||||
|
||||
public:
|
||||
UPROPERTY(BlueprintReadOnly, Category = "Command")
|
||||
TArray<FCommand> Commands;
|
||||
TArray<FRemoteControlActuatorCommand> Commands;
|
||||
|
||||
UPROPERTY(BlueprintReadOnly, Category = "Command")
|
||||
int32 Index = 0;
|
||||
};
|
||||
|
||||
DECLARE_DYNAMIC_MULTICAST_DELEGATE_OneParam(FCommandReady, const FPayload&, Payload);
|
||||
DECLARE_DYNAMIC_MULTICAST_DELEGATE_OneParam(FCommandReady, const FRemoteControlPayload&, Payload);
|
||||
DECLARE_DYNAMIC_MULTICAST_DELEGATE_OneParam(FSocketReady, bool, bSuccessful);
|
||||
|
||||
UCLASS()
|
||||
@ -51,9 +51,9 @@ public:
|
||||
|
||||
ULuckyDataTransferSubsystem();
|
||||
|
||||
virtual void Initialize(FSubsystemCollectionBase& Collection);
|
||||
virtual void Initialize(FSubsystemCollectionBase& Collection) override;
|
||||
|
||||
virtual void Deinitialize();
|
||||
virtual void Deinitialize() override;
|
||||
|
||||
TSharedPtr<IWebSocket> Socket;
|
||||
|
||||
@ -61,7 +61,8 @@ public:
|
||||
|
||||
void Internal_OpenWebsocket(const FString& URL, const FString& Protocol);
|
||||
|
||||
FPayload InterpretData(const FString& Message);
|
||||
// Parse a JSON Payload received from a websocket connection into a Cpp struct
|
||||
static FRemoteControlPayload ParseJsonPayload(const FString& Message);
|
||||
|
||||
UPROPERTY(BlueprintAssignable)
|
||||
FCommandReady OnCommandReady;
|
||||
@ -74,10 +75,10 @@ public:
|
||||
void Callback_OnConnected();
|
||||
|
||||
UFUNCTION()
|
||||
void Callback_OnConnectionError(const FString& Error);
|
||||
void Callback_OnConnectionError(const FString& Error) const;
|
||||
|
||||
UFUNCTION()
|
||||
void Callback_OnMessage(const FString& Message);
|
||||
void Callback_OnMessage(const FString& Message) const;
|
||||
|
||||
UFUNCTION()
|
||||
void Internal_OnMessageSent(const FString& Message);
|
||||
@ -93,7 +94,7 @@ public:
|
||||
void SendMessage(const FString& Message);
|
||||
|
||||
UFUNCTION()
|
||||
void CommandReady(const FPayload& Payload);
|
||||
void CommandReady(const FRemoteControlPayload& Payload) const;
|
||||
|
||||
//---Observations (Sent to server from Unreal)--------------//
|
||||
//Feature Data declarations
|
||||
@ -126,7 +127,7 @@ public:
|
||||
void RegisterSensor(ALuckySensorPawnBase* Sensor);
|
||||
|
||||
UFUNCTION(BlueprintCallable, Meta = (AutoCreateRefTerm = "Path, Comment"), Category = "Capture")
|
||||
bool WriteImageToDisk(const FString& Path, const double Timestamp, const FString& Comment = "");
|
||||
bool WriteImageToDisk(const double Timestamp, const FString& InPath = "", const FString& Comment = "");
|
||||
//-------------------------------------------------------//
|
||||
|
||||
};
|
||||
|
||||
@ -8,11 +8,18 @@ public class LuckyWorldV2 : ModuleRules
|
||||
{
|
||||
PCHUsage = PCHUsageMode.UseExplicitOrSharedPCHs;
|
||||
|
||||
PublicDependencyModuleNames.AddRange(new string[] { "Core", "CoreUObject", "Engine", "InputCore", "EnhancedInput",
|
||||
PublicDependencyModuleNames.AddRange(new string[] {
|
||||
"Core",
|
||||
"CoreUObject",
|
||||
"Engine",
|
||||
"InputCore",
|
||||
"EnhancedInput",
|
||||
"ChaosVehicles",
|
||||
"PhysicsCore",
|
||||
"AsyncLoadingScreen",
|
||||
"BlueprintJson",
|
||||
"Json",
|
||||
"JsonUtilities",
|
||||
"FileHelper",
|
||||
"LuckyMujoco",
|
||||
"LuckyTextWrite",
|
||||
|
||||
@ -7,6 +7,12 @@
|
||||
#include "LuckyDataTransferSubsystem.h"
|
||||
#include "Components/TextRenderComponent.h"
|
||||
#include "Engine/TextRenderActor.h"
|
||||
#include "Misc/FileHelper.h"
|
||||
#include "Misc/Paths.h"
|
||||
#include "Dom/JsonObject.h"
|
||||
#include "Serialization/JsonWriter.h"
|
||||
#include "Serialization/JsonSerializer.h"
|
||||
#include "_Utils/FileUtils.h"
|
||||
|
||||
UEpisodeSubSystem::UEpisodeSubSystem()
|
||||
{
|
||||
@ -16,12 +22,15 @@ UEpisodeSubSystem::UEpisodeSubSystem()
|
||||
void UEpisodeSubSystem::Initialize(FSubsystemCollectionBase& Collection)
|
||||
{
|
||||
Super::Initialize(Collection);
|
||||
if (ULuckyDataTransferSubsystem* DataTransferSubSystem = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
|
||||
{
|
||||
DataTransfer = DataTransferSubSystem;
|
||||
}
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::Deinitialize()
|
||||
{
|
||||
bTickEnabled = false;
|
||||
FTSTicker::GetCoreTicker().RemoveTicker(TickHandle);
|
||||
StopTicking();
|
||||
Super::Deinitialize();
|
||||
}
|
||||
|
||||
@ -45,16 +54,21 @@ void UEpisodeSubSystem::Tick(float DeltaTime)
|
||||
{
|
||||
const bool bIsEpisodeCompleted = CheckEpisodeCompletion();
|
||||
|
||||
if (bIsEpisodeCompleted && CapturedEpisodes <= EpisodesToCapture)
|
||||
// Write Image on the disk
|
||||
if (DataTransfer) DataTransfer->WriteImageToDisk(CurrentRobot->PhysicsSceneProxy->GetMujocoData().time);
|
||||
EpisodeFrames++;
|
||||
|
||||
if (!bIsEpisodeCompleted) return;
|
||||
|
||||
EndEpisode();
|
||||
if (CapturedEpisodes < EpisodesToCapture)
|
||||
{
|
||||
return StartEpisode();
|
||||
StartEpisode();
|
||||
}
|
||||
else
|
||||
{
|
||||
EndTraining();
|
||||
}
|
||||
|
||||
// Here shouldn't we rewrite the frames to know if the episode was a success or a failure?
|
||||
|
||||
// Maybe this should not be done in the tick but after episode completion
|
||||
const auto Payload = CreatePayload();
|
||||
SendEpisodeData(Payload);
|
||||
}
|
||||
}
|
||||
|
||||
@ -68,6 +82,12 @@ void UEpisodeSubSystem::StartTicking()
|
||||
TickHandle = FTSTicker::GetCoreTicker().AddTicker(TickDelegate);
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::StopTicking()
|
||||
{
|
||||
bTickEnabled = false;
|
||||
FTSTicker::GetCoreTicker().RemoveTicker(TickHandle);
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::UpdateDebugTextActor() const
|
||||
{
|
||||
if (!IsValid(DebugTextActor)) return;
|
||||
@ -77,7 +97,7 @@ void UEpisodeSubSystem::UpdateDebugTextActor() const
|
||||
TextRender->SetText(FText::FromString(Txt));
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::StartNewEpisodesSeries(const int32 EpisodesCountIn, FString BaseImageDataPathIn)
|
||||
void UEpisodeSubSystem::StartTraining(const int32 EpisodesCountIn, FString BaseImageDataPathIn, FString TaskDescriptionIn)
|
||||
{
|
||||
// Debug
|
||||
const auto DebugTextActorPtr = UGameplayStatics::GetActorOfClass(this->GetWorld(), ATextRenderActor::StaticClass());
|
||||
@ -97,10 +117,18 @@ void UEpisodeSubSystem::StartNewEpisodesSeries(const int32 EpisodesCountIn, FStr
|
||||
// Data
|
||||
ConfigureDataCapture();
|
||||
BaseImageDataPath = BaseImageDataPathIn;
|
||||
TaskDescription = TaskDescriptionIn;
|
||||
|
||||
StartTicking();
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::EndTraining()
|
||||
{
|
||||
StopTicking();
|
||||
CreateEpisodesStatsJsonFile();
|
||||
// Create jsonl files
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::StartEpisode()
|
||||
{
|
||||
// Robot should be in its ready state - overriden per PilotComponent
|
||||
@ -114,9 +142,6 @@ void UEpisodeSubSystem::StartEpisode()
|
||||
RobotTransform.GetLocation() + RobotTransform.GetRotation().GetForwardVector() * HardCodedRewardDistanceFromRobotPivot * (FMath::RandBool() ? 1 : -1)
|
||||
};
|
||||
|
||||
// DrawDebugLine(this->GetWorld(), EpisodeRewardZone.GetLocation() + FVector::UpVector * 70, EpisodeRewardZone.GetLocation(), FColor::Red, true);
|
||||
// DrawDebugLine(this->GetWorld(), RobotTransform.GetLocation() + FVector::UpVector * 70, RobotTransform.GetLocation(), FColor::Blue, true);
|
||||
|
||||
// Ask the bot to give a reachable location for the Training Object Transform
|
||||
EpisodeObjectBaseTransform = CurrentRobot->RobotPilotComponent->GetReachableTransform();
|
||||
|
||||
@ -134,6 +159,24 @@ void UEpisodeSubSystem::StartEpisode()
|
||||
UpdateDebugTextActor();
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::EndEpisode()
|
||||
{
|
||||
// Gather the robot data
|
||||
const FTrainingEpisodeData TrainingEpisodeData = CurrentRobot->RobotPilotComponent->GetTrainingEpisodeData();
|
||||
|
||||
// Create episodes_stats.jsonl single line and append to EpisodeStatLines
|
||||
CreateEpisodeStatJsonLine(TrainingEpisodeData);
|
||||
|
||||
// create a parquet file
|
||||
CreateEpisodeParquetFile();
|
||||
|
||||
// Convert images into video
|
||||
// TODO Find a good FFMPEG plugin - maybe the Unreal base one is good
|
||||
|
||||
// Reset values for the next episode
|
||||
EpisodeFrames = 0;
|
||||
}
|
||||
|
||||
bool UEpisodeSubSystem::CheckEpisodeCompletion()
|
||||
{
|
||||
const auto GeomTransform = CurrentRobot->PhysicsSceneProxy->GetGeometryTransform(EpisodeTargetObject->MainActorBody.GetName());
|
||||
@ -184,53 +227,94 @@ void UEpisodeSubSystem::FindRobotPawnFromScene()
|
||||
}
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::InitCameras()
|
||||
{
|
||||
// TODO Fix the spawning of sensors in Cpp and spawn them using a config?
|
||||
// TODO How people can move the camera themselves?
|
||||
|
||||
// Find all sensors in the scene
|
||||
TArray<AActor*> Sensors;
|
||||
UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), ALuckySensorPawnBase::StaticClass(), Sensors);
|
||||
|
||||
for (const auto Sensor : Sensors)
|
||||
{
|
||||
if (const auto Camera = Cast<ALuckySensorPawnBase>(Sensor)) Cameras.Add(Camera);
|
||||
}
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::ConfigureDataCapture()
|
||||
{
|
||||
if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
|
||||
if (!DataTransfer) return;
|
||||
DataTransfer->CreateCaptureSessionID();
|
||||
InitCameras();
|
||||
for (const auto& Cam : Cameras)
|
||||
{
|
||||
//Do this before your tick operation - shouldn't happen on tick
|
||||
//Connect to websocket and create session id
|
||||
DataTransfer->ConnectToWebsocket("ws://127.0.0.1:3000", "");
|
||||
DataTransfer->CreateCaptureSessionID();
|
||||
DataTransfer->RegisterSensor(Cam.Get());
|
||||
Cam->SensorInfo.bActive = true;
|
||||
}
|
||||
}
|
||||
|
||||
FObservationPayload UEpisodeSubSystem::CreatePayload()
|
||||
void UEpisodeSubSystem::CreateEpisodeStatJsonLine(const FTrainingEpisodeData& TrainingEpisodeData)
|
||||
{
|
||||
// CurrentRobot->Cameras
|
||||
// CurrentRobot -> Tell JB what he should expose on the RobotPawn
|
||||
// const auto TimeStamp = CurrentRobot->PhysicsSceneProxy->GetMujocoData().time;
|
||||
// const auto So100PilotCmp = Cast<URobotPilotSO100Component>(CurrentRobot->RobotPilotComponent);
|
||||
// const auto Joints = So100PilotCmp->GetCurrentControlsFromPhysicScene();
|
||||
// EpisodeStatLines.
|
||||
const TSharedPtr<FJsonObject> Root = MakeShared<FJsonObject>();
|
||||
Root->SetNumberField("episode_index", CapturedEpisodes);
|
||||
|
||||
// Tick operation
|
||||
// Create the payload
|
||||
return FObservationPayload {
|
||||
// timestamp goes here - FString,
|
||||
// "observation", //just leave this because this is what ethan and anuj will expect
|
||||
// enter a message here - FString,
|
||||
// TMap of FString (Actuator name or index), and Float (value of actuator)
|
||||
// Camera info struct goes here, don't worry about this for now, just use TArray<FObservationCameraObject>()
|
||||
// What about episode success?
|
||||
// How to invalidate data
|
||||
};
|
||||
const TSharedPtr<FJsonObject> Stats = MakeShared<FJsonObject>();
|
||||
Stats->SetObjectField("action", MakeShared<FJsonObject>(TrainingEpisodeData.ControlsStats));
|
||||
Stats->SetObjectField("observation.state", MakeShared<FJsonObject>(TrainingEpisodeData.JointsStats));
|
||||
|
||||
// TODO Once all json and parquet files are written on disk and the PR is merged into main and tested, we will do it
|
||||
// TODO "observation.images.webcam"
|
||||
// TODO "timestamp"
|
||||
// TODO "frame_index"
|
||||
// TODO "episode_index"
|
||||
// TODO "index"
|
||||
// TODO "task_index"
|
||||
|
||||
// Append
|
||||
Root->SetObjectField("stats", Stats);
|
||||
|
||||
// Serialize into FString
|
||||
FString Output;
|
||||
const TSharedRef< TJsonWriter< TCHAR, TCondensedJsonPrintPolicy<TCHAR> > > Writer = TJsonWriterFactory< TCHAR, TCondensedJsonPrintPolicy<TCHAR> >::Create(&Output);
|
||||
FJsonSerializer::Serialize(Root.ToSharedRef(), Writer);
|
||||
EpisodeStatLines.Add(Output);
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::SendEpisodeData(const FObservationPayload& Payload) const
|
||||
void UEpisodeSubSystem::CreateEpisodeParquetFile()
|
||||
{
|
||||
if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
|
||||
{
|
||||
// Here generate the path for each image?
|
||||
// DataTransfer->WriteImageToDisk(BaseImageDataPath, 0.f);
|
||||
|
||||
// Don't send data if socket is disconnected
|
||||
if (!DataTransfer->Socket->IsConnected()) return;
|
||||
|
||||
// Send the Data
|
||||
//Queue and convert the payload to json
|
||||
DataTransfer->CreateJsonPayload_Observation(Payload);
|
||||
|
||||
//Send the payload over websocket
|
||||
DataTransfer->SendMessage(DataTransfer->ObservationPayloadString);
|
||||
}
|
||||
// TODO Use Anuj plugin to create one parquet file per episode
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::ConvertImagesToVideo()
|
||||
{
|
||||
// TODO Once every json and parquet tasks are done
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::CreateEpisodesStatsJsonFile()
|
||||
{
|
||||
// TODO Do not use FJsonObject - simply concat the FStrings into a file
|
||||
UFileUtils::WriteJsonlFile(EpisodeStatLines, FPaths::ProjectSavedDir(), FString("episodes_stats"));
|
||||
|
||||
// Create a jsonl file and store in the correct directory
|
||||
// concat TArray<FString> EpisodeStatLines into a single file
|
||||
// https://huggingface.co/datasets/youliangtan/so100_strawberry_grape/blob/main/meta/episodes_stats.jsonl
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::CreateEpisodesJsonFile()
|
||||
{
|
||||
// Create a jsonl file and store in the correct directory
|
||||
// https://huggingface.co/datasets/youliangtan/so100_strawberry_grape/blob/main/meta/episodes.jsonl
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::CreateInfoJsonFile()
|
||||
{
|
||||
// https://huggingface.co/datasets/youliangtan/so100_strawberry_grape/blob/main/meta/info.json
|
||||
}
|
||||
|
||||
void UEpisodeSubSystem::CreateTasksJsonFile()
|
||||
{
|
||||
// https://huggingface.co/datasets/youliangtan/so100_strawberry_grape/blob/main/meta/tasks.jsonl
|
||||
}
|
||||
|
||||
|
||||
@ -60,3 +60,23 @@ void URobotPilotComponent::SetRobotTarget(const FTransform& TargetTransformIn)
|
||||
void URobotPilotComponent::SetRobotCurrentRewardZone(const FTransform& RewardTransformIn)
|
||||
{
|
||||
}
|
||||
|
||||
void URobotPilotComponent::ReceiveRemoteCommand(const FRemoteControlPayload& RemoteRobotPayload)
|
||||
{
|
||||
}
|
||||
|
||||
FJsonObject URobotPilotComponent::GetBufferedControlsData()
|
||||
{
|
||||
return FJsonObject();
|
||||
}
|
||||
|
||||
FJsonObject URobotPilotComponent::GetBufferedJointsData()
|
||||
{
|
||||
return FJsonObject();
|
||||
}
|
||||
|
||||
FTrainingEpisodeData URobotPilotComponent::GetTrainingEpisodeData()
|
||||
{
|
||||
return FTrainingEpisodeData();
|
||||
}
|
||||
|
||||
|
||||
@ -1,4 +1,6 @@
|
||||
#include "Robot/PilotComponent/RobotPilotSO100Component.h"
|
||||
|
||||
#include "LuckyDataTransferSubsystem.h"
|
||||
#include "Actors/MujocoVolumeActor.h"
|
||||
#include "Kismet/KismetMathLibrary.h"
|
||||
#include "Robot/RobotPawn.h"
|
||||
@ -22,10 +24,9 @@ void URobotPilotSO100Component::BeginPlay()
|
||||
void URobotPilotSO100Component::TickComponent(float DeltaTime, enum ELevelTick TickType,
|
||||
FActorComponentTickFunction* ThisTickFunction)
|
||||
{
|
||||
// Super::TickComponent(DeltaTime, TickType, ThisTickFunction);
|
||||
// const auto Joints = GetCurrentJointsFromPhysicsScene();
|
||||
// const auto Controls = GetCurrentControlsFromPhysicScene();
|
||||
// UE_LOG(LogTemp, Log, TEXT("Joint: %f - Control: %f - Delta: %f"), Joints.Jaw, Controls.Jaw, Controls.Jaw - Joints.Jaw);
|
||||
// We can buffer data in the tick because it will never be called in parallel of MujocoVolumeActor tick
|
||||
ControlsDataBuffer.Add(GetCurrentControlsFromPhysicScene());
|
||||
JointsDataBuffer.Add(GetCurrentJointsFromPhysicsScene());
|
||||
}
|
||||
|
||||
FTransform URobotPilotSO100Component::GetReachableTransform()
|
||||
@ -42,8 +43,6 @@ FTransform URobotPilotSO100Component::GetReachableTransform()
|
||||
const float RandomYaw = MaxYaw * FMath::RandRange(0., 1.) * (FMath::RandBool() ? 1 : -1);
|
||||
const FQuat RandomRotation = FRotator{0,RandomYaw,0}.Quaternion() * ArmWorldRotation;
|
||||
|
||||
|
||||
|
||||
// Compute Random Range within reach of the arm and add this to pivot location
|
||||
// Add a bit more than the Jaw Offset - TODO Offsets must be better computed
|
||||
const float RandomRange = JawOffset.X + MaxRange * FMath::RandRange(0.1f, 1.f);
|
||||
@ -54,11 +53,6 @@ FTransform URobotPilotSO100Component::GetReachableTransform()
|
||||
RewardAxis.Z = 0; // Nullify Z to keep a 2D vector -> ensure the geometry roll/pitch are 0
|
||||
const FRotator TowardPivotRotation = UKismetMathLibrary::MakeRotFromXZ(RewardAxis, FVector::UpVector);
|
||||
|
||||
// Debug
|
||||
// DrawDebugLine(this->GetWorld(), ArmPivotLocation + ArmWorldRotation.GetForwardVector() * 70, ArmPivotLocation, FColor::Green, true);
|
||||
// DrawDebugLine(this->GetWorld(), ArmPivotLocation + FVector::UpVector * 70, ArmPivotLocation, FColor::Red, true);
|
||||
// DrawDebugLine(this->GetWorld(), RandomLocation, RandomLocation + TowardPivotRotation.Quaternion().GetForwardVector() * -50 , FColor::Blue, true);
|
||||
|
||||
// Return the Object Transform
|
||||
return FTransform(TowardPivotRotation, RandomLocation);
|
||||
}
|
||||
@ -90,6 +84,182 @@ void URobotPilotSO100Component::SetRobotCurrentRewardZone(const FTransform& Rewa
|
||||
bDropZoneIsRight = FVector::DotProduct(RobotOwner->RobotActor->GetActorRotation().Quaternion().GetRightVector(), DirectionToTarget) > 0.f;
|
||||
}
|
||||
|
||||
void URobotPilotSO100Component::ReceiveRemoteCommand(const FRemoteControlPayload& RemoteRobotPayload)
|
||||
{
|
||||
for (const auto& [ActuatorName, ActuatorValue] : RemoteRobotPayload.Commands)
|
||||
{
|
||||
// Will print an error message if it doesn't exist
|
||||
RobotOwner->PhysicsSceneProxy->SetActuatorValue(ActuatorName, ActuatorValue);
|
||||
}
|
||||
}
|
||||
|
||||
FTrainingEpisodeData URobotPilotSO100Component::GetTrainingEpisodeData()
|
||||
{
|
||||
const auto TrainingData = FTrainingEpisodeData
|
||||
{
|
||||
GetBufferedJointsData(),
|
||||
GetBufferedControlsData(),
|
||||
GetStats(ControlsDataBuffer),
|
||||
GetStats(JointsDataBuffer)
|
||||
};
|
||||
|
||||
ControlsDataBuffer.Empty();
|
||||
JointsDataBuffer.Empty();
|
||||
return TrainingData;
|
||||
}
|
||||
|
||||
FJsonObject URobotPilotSO100Component::GetBufferedControlsData()
|
||||
{
|
||||
auto BufferedData = FJsonObject();
|
||||
TArray<TSharedPtr<FJsonValue>> ControlsArray;
|
||||
|
||||
for (const auto& Control : ControlsDataBuffer)
|
||||
{
|
||||
TSharedPtr<FJsonObject> ControlJson = MakeShared<FJsonObject>();
|
||||
|
||||
ControlJson->SetNumberField(TEXT("Rotation"), Control.Rotation);
|
||||
ControlJson->SetNumberField(TEXT("Pitch"), Control.Pitch);
|
||||
ControlJson->SetNumberField(TEXT("Elbow"), Control.Elbow);
|
||||
ControlJson->SetNumberField(TEXT("WristPitch"), Control.WristPitch);
|
||||
ControlJson->SetNumberField(TEXT("WristRoll"), Control.WristRoll);
|
||||
ControlJson->SetNumberField(TEXT("Jaw"), Control.Jaw);
|
||||
|
||||
ControlsArray.Add(MakeShared<FJsonValueObject>(ControlJson));
|
||||
}
|
||||
|
||||
BufferedData.SetArrayField(TEXT("action"), ControlsArray);
|
||||
|
||||
return BufferedData;
|
||||
}
|
||||
|
||||
FJsonObject URobotPilotSO100Component::GetBufferedJointsData()
|
||||
{
|
||||
auto BufferedData = FJsonObject();
|
||||
TArray<TSharedPtr<FJsonValue>> ControlsArray;
|
||||
|
||||
for (const auto& Joint : JointsDataBuffer)
|
||||
{
|
||||
TSharedPtr<FJsonObject> ControlJson = MakeShared<FJsonObject>();
|
||||
|
||||
ControlJson->SetNumberField(TEXT("Rotation"), Joint.Rotation);
|
||||
ControlJson->SetNumberField(TEXT("Pitch"), Joint.Pitch);
|
||||
ControlJson->SetNumberField(TEXT("Elbow"), Joint.Elbow);
|
||||
ControlJson->SetNumberField(TEXT("WristPitch"), Joint.WristPitch);
|
||||
ControlJson->SetNumberField(TEXT("WristRoll"), Joint.WristRoll);
|
||||
ControlJson->SetNumberField(TEXT("Jaw"), Joint.Jaw);
|
||||
|
||||
ControlsArray.Add(MakeShared<FJsonValueObject>(ControlJson));
|
||||
}
|
||||
|
||||
BufferedData.SetArrayField(TEXT("observation.state"), ControlsArray);
|
||||
|
||||
return BufferedData;
|
||||
}
|
||||
|
||||
FJsonObject URobotPilotSO100Component::GetStats(const TArray<FSo100Actuators>& ActuatorStates)
|
||||
{
|
||||
FSo100Actuators Min;
|
||||
FSo100Actuators Max;
|
||||
FSo100Actuators Sum = FSo100Actuators();
|
||||
FSo100Actuators Variance = FSo100Actuators();
|
||||
FSo100Actuators Mean;
|
||||
FSo100Actuators StdDev;
|
||||
|
||||
// Compute Sum
|
||||
for (const auto& State : ActuatorStates)
|
||||
{
|
||||
Sum.Rotation += State.Rotation;
|
||||
Sum.Pitch += State.Pitch;
|
||||
Sum.Elbow += State.Elbow;
|
||||
Sum.WristPitch += State.WristPitch;
|
||||
Sum.WristRoll += State.WristRoll;
|
||||
Sum.Jaw += State.Jaw;
|
||||
|
||||
// Assign Max
|
||||
if(State.Rotation > Max.Rotation) Max.Rotation = State.Rotation;
|
||||
if(State.Pitch > Max.Pitch) Max.Pitch = State.Pitch;
|
||||
if(State.Elbow > Max.Elbow) Max.Elbow = State.Elbow;
|
||||
if(State.WristPitch > Max.WristPitch) Max.WristPitch = State.WristPitch;
|
||||
if(State.WristRoll > Max.WristRoll) Max.WristRoll = State.WristRoll;
|
||||
if(State.Jaw > Max.Jaw) Max.Jaw = State.Jaw;
|
||||
|
||||
// Assing min
|
||||
if(State.Rotation < Min.Rotation) Min.Rotation = State.Rotation;
|
||||
if(State.Pitch < Min.Pitch) Min.Pitch = State.Pitch;
|
||||
if(State.Elbow < Min.Elbow) Min.Elbow = State.Elbow;
|
||||
if(State.WristPitch < Min.WristPitch) Min.WristPitch = State.WristPitch;
|
||||
if(State.WristRoll < Min.WristRoll) Min.WristRoll = State.WristRoll;
|
||||
if(State.Jaw < Min.Jaw) Min.Jaw = State.Jaw;
|
||||
|
||||
}
|
||||
|
||||
// Compute Mean
|
||||
Mean.Rotation = Sum.Rotation / ActuatorStates.Num();
|
||||
Mean.Pitch = Sum.Pitch / ActuatorStates.Num();
|
||||
Mean.Elbow = Sum.Elbow / ActuatorStates.Num();
|
||||
Mean.WristPitch = Sum.WristPitch / ActuatorStates.Num();
|
||||
Mean.WristRoll = Sum.WristRoll / ActuatorStates.Num();
|
||||
Mean.Jaw = Sum.Jaw / ActuatorStates.Num();
|
||||
|
||||
// Compute Variance
|
||||
// Pre step
|
||||
for (const auto& State : ActuatorStates)
|
||||
{
|
||||
const double DiffRotation = State.Rotation - Mean.Rotation;
|
||||
const double DiffPitch = State.Pitch - Mean.Pitch;
|
||||
const double DiffElbow = State.Elbow - Mean.Elbow;
|
||||
const double DiffWristPitch = State.WristPitch - Mean.WristPitch;
|
||||
const double DiffWristRoll = State.WristRoll - Mean.WristRoll;
|
||||
const double DiffJaw = State.Jaw - Mean.Jaw;
|
||||
|
||||
Variance.Rotation += DiffRotation * DiffRotation;
|
||||
Variance.Pitch += DiffPitch * DiffPitch;
|
||||
Variance.Elbow += DiffElbow * DiffElbow;
|
||||
Variance.WristPitch += DiffWristPitch * DiffWristPitch;
|
||||
Variance.WristRoll += DiffWristRoll * DiffWristRoll;
|
||||
Variance.Jaw += DiffJaw * DiffJaw;
|
||||
}
|
||||
|
||||
// Final Variance Value per actuator
|
||||
Variance.Rotation /= ActuatorStates.Num();
|
||||
Variance.Pitch /= ActuatorStates.Num();
|
||||
Variance.Elbow /= ActuatorStates.Num();
|
||||
Variance.WristPitch /= ActuatorStates.Num();
|
||||
Variance.WristRoll /= ActuatorStates.Num();
|
||||
Variance.Jaw /= ActuatorStates.Num();
|
||||
|
||||
// Standard Deviation
|
||||
StdDev.Rotation = FMath::Sqrt(Variance.Rotation);
|
||||
StdDev.Pitch = FMath::Sqrt(Variance.Pitch);
|
||||
StdDev.Elbow = FMath::Sqrt(Variance.Elbow);
|
||||
StdDev.WristPitch = FMath::Sqrt(Variance.WristPitch);
|
||||
StdDev.WristRoll = FMath::Sqrt(Variance.WristRoll);
|
||||
StdDev.Jaw = FMath::Sqrt(Variance.Jaw);
|
||||
|
||||
// Here create a json in the following format
|
||||
FJsonObject StatsJson = FJsonObject();
|
||||
|
||||
// Helper lambda to convert FSo100Actuators to TArray<TSharedPtr<FJsonValue>>
|
||||
auto ActuatorsToJsonArray = [](const FSo100Actuators& Actuators) -> TArray<TSharedPtr<FJsonValue>>
|
||||
{
|
||||
return {
|
||||
MakeShared<FJsonValueNumber>(Actuators.Rotation),
|
||||
MakeShared<FJsonValueNumber>(Actuators.Pitch),
|
||||
MakeShared<FJsonValueNumber>(Actuators.Elbow),
|
||||
MakeShared<FJsonValueNumber>(Actuators.WristPitch),
|
||||
MakeShared<FJsonValueNumber>(Actuators.WristRoll),
|
||||
MakeShared<FJsonValueNumber>(Actuators.Jaw)
|
||||
};
|
||||
};
|
||||
|
||||
StatsJson.SetArrayField(TEXT("min"), ActuatorsToJsonArray(Min));
|
||||
StatsJson.SetArrayField(TEXT("max"), ActuatorsToJsonArray(Max));
|
||||
StatsJson.SetArrayField(TEXT("mean"), ActuatorsToJsonArray(Mean));
|
||||
StatsJson.SetArrayField(TEXT("std"), ActuatorsToJsonArray(StdDev));
|
||||
StatsJson.SetArrayField(TEXT("count"), { MakeShared<FJsonValueNumber>(ActuatorStates.Num()) });
|
||||
|
||||
return StatsJson;
|
||||
}
|
||||
|
||||
void URobotPilotSO100Component::PrintCurrentActuators() const
|
||||
{
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
#include "Robot/RobotPawn.h"
|
||||
|
||||
#include "LuckyDataTransferSubsystem.h"
|
||||
#include "LuckySensorPawnBase.h"
|
||||
#include "Kismet/GameplayStatics.h"
|
||||
#include "Kismet/KismetMathLibrary.h"
|
||||
@ -20,7 +21,6 @@ void ARobotPawn::BeginPlay()
|
||||
void ARobotPawn::InitRobot()
|
||||
{
|
||||
InitPilotComponent();
|
||||
InitCamera();
|
||||
}
|
||||
|
||||
void ARobotPawn::InitPilotComponent()
|
||||
@ -70,17 +70,19 @@ void ARobotPawn::InitPilotComponent()
|
||||
}
|
||||
}
|
||||
|
||||
void ARobotPawn::InitCamera()
|
||||
void ARobotPawn::EnableRemoteControl()
|
||||
{
|
||||
// TODO Fix the spawning of sensors in Cpp and spawn them using a config?
|
||||
// TODO How people can move the camera themselves?
|
||||
|
||||
// Find all sensors in the scene
|
||||
TArray<AActor*> Sensors;
|
||||
UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), ALuckySensorPawnBase::StaticClass(), Sensors);
|
||||
|
||||
for (const auto Sensor : Sensors)
|
||||
// Get subsystem
|
||||
if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
|
||||
{
|
||||
if (const auto Camera = Cast<ALuckySensorPawnBase>(Sensor)) Cameras.Add(Camera);
|
||||
// Connect first if necessary
|
||||
if (!DataTransfer->Socket->IsConnected())
|
||||
{
|
||||
DataTransfer->ConnectToWebsocket("ws://127.0.0.1:3000", "");
|
||||
}
|
||||
|
||||
// TODO Should we wait for connection to be successful before binding OnCommandReady?
|
||||
DataTransfer->OnCommandReady.AddDynamic(this->RobotPilotComponent, &URobotPilotComponent::ReceiveRemoteCommand);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
24
Source/LuckyWorldV2/Private/_Utils/FileUtils.cpp
Normal file
24
Source/LuckyWorldV2/Private/_Utils/FileUtils.cpp
Normal file
@ -0,0 +1,24 @@
|
||||
#include "_Utils/FileUtils.h"
|
||||
#include "Misc/FileHelper.h"
|
||||
#include "Misc/Paths.h"
|
||||
|
||||
UFileUtils::UFileUtils()
|
||||
{
|
||||
}
|
||||
|
||||
|
||||
|
||||
bool UFileUtils::WriteJsonlFile(const TArray<FString>& JsonLines, const FString& BasePath, const FString& FileName)
|
||||
{
|
||||
// Ensure the directory exists
|
||||
IFileManager::Get().MakeDirectory(*BasePath, true);
|
||||
|
||||
// Construct the full file path
|
||||
const FString FullFilePath = FPaths::Combine(BasePath, FileName + TEXT(".jsonl"));
|
||||
|
||||
// Join the array into one string with line breaks
|
||||
const FString FileContent = FString::Join(JsonLines, TEXT(""));
|
||||
|
||||
// Write to file
|
||||
return FFileHelper::SaveStringToFile(FileContent, *FullFilePath);
|
||||
}
|
||||
@ -1,15 +1,34 @@
|
||||
#pragma once
|
||||
#include "CoreMinimal.h"
|
||||
#include "ObservationData.h"
|
||||
#include "Robot/PilotComponent/RobotPilotComponent.h"
|
||||
#include "Subsystems/WorldSubsystem.h"
|
||||
#include "Stats/Stats.h"
|
||||
#include "EpisodeSubSystem.generated.h"
|
||||
|
||||
|
||||
class ULuckyDataTransferSubsystem;
|
||||
class ALuckySensorPawnBase;
|
||||
class ATextRenderActor;
|
||||
class AMujocoStaticMeshActor;
|
||||
class ARobotPawn;
|
||||
|
||||
USTRUCT()
|
||||
struct FTrainingEpisode
|
||||
{
|
||||
GENERATED_BODY()
|
||||
UPROPERTY()
|
||||
int32 EpisodeIndex = -1;
|
||||
|
||||
UPROPERTY()
|
||||
TArray<FString> Tasks;
|
||||
|
||||
UPROPERTY()
|
||||
int32 Length = -1;
|
||||
};
|
||||
|
||||
|
||||
|
||||
|
||||
UCLASS()
|
||||
class LUCKYWORLDV2_API UEpisodeSubSystem : public UWorldSubsystem
|
||||
{
|
||||
@ -29,6 +48,7 @@ public:
|
||||
// It will allows us to remove all the episode logic from the SubSystem and having different types of episodes
|
||||
void Tick(float DeltaTime);
|
||||
void StartTicking();
|
||||
void StopTicking();
|
||||
FTSTicker::FDelegateHandle TickHandle;
|
||||
bool bTickEnabled = true;
|
||||
|
||||
@ -50,8 +70,11 @@ public:
|
||||
* Called by the UI when pressing the "Capture" button
|
||||
*/
|
||||
UFUNCTION(BlueprintCallable)
|
||||
void StartNewEpisodesSeries(int32 EpisodesCountIn, FString BaseImageDataPathIn);
|
||||
void StartTraining(int32 EpisodesCountIn, FString BaseImageDataPathIn, FString TaskDescriptionIn);
|
||||
void EndTraining();
|
||||
|
||||
UPROPERTY()
|
||||
ULuckyDataTransferSubsystem* DataTransfer = nullptr;
|
||||
|
||||
|
||||
private:
|
||||
@ -59,6 +82,7 @@ private:
|
||||
// ------- FLOW --------
|
||||
// ---------------------
|
||||
void StartEpisode();
|
||||
void EndEpisode();
|
||||
FTransform EpisodeRewardZone = FTransform::Identity;
|
||||
float EpisodeRewardZoneRadius = 5.f; // TODO Not hardcode it - or only in the Robot? - Maybe we want different scenarios for the robot
|
||||
bool CheckEpisodeCompletion();
|
||||
@ -85,24 +109,32 @@ private:
|
||||
UPROPERTY()
|
||||
TObjectPtr<ARobotPawn> CurrentRobot;
|
||||
|
||||
|
||||
|
||||
// ---------------------
|
||||
// ------ SENSORS ------
|
||||
// ---------------------
|
||||
void InitCameras();
|
||||
TArray<TObjectPtr<ALuckySensorPawnBase>> Cameras;
|
||||
|
||||
// --------------------
|
||||
// ------- DATA -------
|
||||
// --------------------
|
||||
FString BaseImageDataPath;
|
||||
|
||||
FString TaskDescription;
|
||||
int32 EpisodeFrames = 0;
|
||||
// Noah here add anything you need
|
||||
void ConfigureDataCapture();
|
||||
|
||||
FObservationPayload CreatePayload();
|
||||
|
||||
void SendEpisodeData(const FObservationPayload& Payload) const;
|
||||
|
||||
|
||||
|
||||
// End of episode tasks
|
||||
void CreateEpisodeStatJsonLine(const FTrainingEpisodeData& TrainingEpisodeData);
|
||||
void CreateEpisodeParquetFile();
|
||||
void ConvertImagesToVideo();
|
||||
|
||||
// End of training files
|
||||
TArray<FString> EpisodeStatLines;
|
||||
void CreateEpisodesStatsJsonFile();
|
||||
void CreateEpisodesJsonFile();
|
||||
void CreateInfoJsonFile();
|
||||
void CreateTasksJsonFile();
|
||||
};
|
||||
|
||||
|
||||
|
||||
@ -2,6 +2,8 @@
|
||||
#include "CoreMinimal.h"
|
||||
#include "RobotPilotComponent.generated.h"
|
||||
|
||||
struct FRemoteControlPayload;
|
||||
|
||||
USTRUCT(BlueprintType)
|
||||
struct FRobotActuators
|
||||
{
|
||||
@ -10,6 +12,17 @@ struct FRobotActuators
|
||||
// What will be in common?
|
||||
};
|
||||
|
||||
USTRUCT(BlueprintType)
|
||||
struct FTrainingEpisodeData
|
||||
{
|
||||
GENERATED_BODY()
|
||||
FJsonObject JointsStates; // The total series of joint values per frame
|
||||
FJsonObject ControlsStates; // The total series of joint values per frame
|
||||
FJsonObject JointsStats; // The min / max / mean / std for joints series
|
||||
FJsonObject ControlsStats; // The min / max / mean / std for controls series
|
||||
};
|
||||
|
||||
|
||||
class ARobotPawn;
|
||||
|
||||
UCLASS(Blueprintable)
|
||||
@ -34,6 +47,15 @@ public:
|
||||
virtual void SetRobotTarget(const FTransform& TargetTransformIn);
|
||||
virtual void SetRobotCurrentRewardZone(const FTransform& RewardTransformIn);
|
||||
|
||||
UFUNCTION()
|
||||
virtual void ReceiveRemoteCommand(const FRemoteControlPayload& RemoteRobotPayload);
|
||||
|
||||
// Data
|
||||
virtual FJsonObject GetBufferedControlsData();
|
||||
virtual FJsonObject GetBufferedJointsData();
|
||||
virtual FTrainingEpisodeData GetTrainingEpisodeData();
|
||||
|
||||
|
||||
|
||||
protected: // Child class need access
|
||||
// Only to easy access within the component
|
||||
|
||||
@ -3,6 +3,8 @@
|
||||
#include "Robot/PilotComponent/RobotPilotComponent.h"
|
||||
#include "RobotPilotSO100Component.generated.h"
|
||||
|
||||
struct FRemoteControlPayload;
|
||||
|
||||
USTRUCT(BlueprintType)
|
||||
struct FSo100Actuators
|
||||
{
|
||||
@ -42,7 +44,17 @@ public:
|
||||
|
||||
virtual void SetRobotTarget(const FTransform& TargetTransformIn) override;
|
||||
virtual void SetRobotCurrentRewardZone(const FTransform& RewardTransformIn) override;
|
||||
virtual void ReceiveRemoteCommand(const FRemoteControlPayload& RemoteRobotPayload) override;
|
||||
|
||||
// ------------------
|
||||
// ------ DATA ------
|
||||
// ------------------
|
||||
virtual FTrainingEpisodeData GetTrainingEpisodeData() override;
|
||||
virtual FJsonObject GetBufferedControlsData() override;
|
||||
virtual FJsonObject GetBufferedJointsData() override;
|
||||
TArray<FSo100Actuators> ControlsDataBuffer;
|
||||
TArray<FSo100Actuators> JointsDataBuffer;
|
||||
static FJsonObject GetStats(const TArray<FSo100Actuators>& ActuatorStates);
|
||||
|
||||
private:
|
||||
FTransform TargetTransform;
|
||||
|
||||
@ -41,10 +41,5 @@ public:
|
||||
UFUNCTION(BlueprintCallable)
|
||||
void InitPilotComponent(); // This should have Robot type as parameter?
|
||||
|
||||
// ---------------------
|
||||
// ------ SENSORS ------
|
||||
// ---------------------
|
||||
void InitCamera();
|
||||
UPROPERTY(EditAnywhere, BlueprintReadWrite)
|
||||
TArray<TObjectPtr<ALuckySensorPawnBase>> Cameras;
|
||||
void EnableRemoteControl();
|
||||
};
|
||||
|
||||
14
Source/LuckyWorldV2/Public/_Utils/FileUtils.h
Normal file
14
Source/LuckyWorldV2/Public/_Utils/FileUtils.h
Normal file
@ -0,0 +1,14 @@
|
||||
#pragma once
|
||||
#include "CoreMinimal.h"
|
||||
#include "FileUtils.generated.h"
|
||||
|
||||
UCLASS()
|
||||
class LUCKYWORLDV2_API UFileUtils : public UObject
|
||||
{
|
||||
GENERATED_BODY()
|
||||
|
||||
public:
|
||||
UFileUtils();
|
||||
|
||||
static bool WriteJsonlFile(const TArray<FString>& JsonLines, const FString& BasePath, const FString& FileName);
|
||||
};
|
||||
Reference in New Issue
Block a user