diff --git a/Content/Blueprint/RobotPawnActors/BP_mujokoSO_100.uasset b/Content/Blueprint/RobotPawnActors/BP_mujokoSO_100.uasset index c208af64..87a704d6 100644 Binary files a/Content/Blueprint/RobotPawnActors/BP_mujokoSO_100.uasset and b/Content/Blueprint/RobotPawnActors/BP_mujokoSO_100.uasset differ diff --git a/Content/Developers/Wdev/Robots/BP_SoArm100robot.uasset b/Content/Developers/Wdev/Robots/BP_SoArm100robot.uasset index ffa0ac0e..ba52d763 100644 Binary files a/Content/Developers/Wdev/Robots/BP_SoArm100robot.uasset and b/Content/Developers/Wdev/Robots/BP_SoArm100robot.uasset differ diff --git a/Source/LuckyWorldV2/LuckyWorldV2.Build.cs b/Source/LuckyWorldV2/LuckyWorldV2.Build.cs index e7a8cf2f..8a50f741 100644 --- a/Source/LuckyWorldV2/LuckyWorldV2.Build.cs +++ b/Source/LuckyWorldV2/LuckyWorldV2.Build.cs @@ -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", diff --git a/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp b/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp index b313574d..c5848e5a 100644 --- a/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp +++ b/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp @@ -7,6 +7,11 @@ #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" UEpisodeSubSystem::UEpisodeSubSystem() { @@ -16,12 +21,15 @@ UEpisodeSubSystem::UEpisodeSubSystem() void UEpisodeSubSystem::Initialize(FSubsystemCollectionBase& Collection) { Super::Initialize(Collection); + if (ULuckyDataTransferSubsystem* DataTransferSubSystem = GetWorld()->GetSubsystem()) + { + DataTransfer = DataTransferSubSystem; + } } void UEpisodeSubSystem::Deinitialize() { - bTickEnabled = false; - FTSTicker::GetCoreTicker().RemoveTicker(TickHandle); + StopTicking(); Super::Deinitialize(); } @@ -45,16 +53,19 @@ void UEpisodeSubSystem::Tick(float DeltaTime) { const bool bIsEpisodeCompleted = CheckEpisodeCompletion(); + // Write Image on the disk + if (DataTransfer) DataTransfer->WriteImageToDisk(CurrentRobot->PhysicsSceneProxy->GetMujocoData().time); + EpisodeFrames++; + if (bIsEpisodeCompleted && CapturedEpisodes <= EpisodesToCapture) { - return StartEpisode(); + EndEpisode(); + 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 +79,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 +94,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 +114,17 @@ void UEpisodeSubSystem::StartNewEpisodesSeries(const int32 EpisodesCountIn, FStr // Data ConfigureDataCapture(); BaseImageDataPath = BaseImageDataPathIn; + TaskDescription = TaskDescriptionIn; StartTicking(); } +void UEpisodeSubSystem::EndTraining() +{ + StopTicking(); + // Create jsonl files +} + void UEpisodeSubSystem::StartEpisode() { // Robot should be in its ready state - overriden per PilotComponent @@ -134,6 +158,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()); @@ -201,64 +243,76 @@ void UEpisodeSubSystem::InitCameras() void UEpisodeSubSystem::ConfigureDataCapture() { - if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem()) + 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(); - - InitCameras(); - for (const auto& Cam : Cameras) - { - DataTransfer->RegisterSensor(Cam.Get()); - Cam->SensorInfo.bActive = true; - } + 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(CurrentRobot->RobotPilotComponent); - // const auto Joints = So100PilotCmp->GetCurrentControlsFromPhysicScene(); + // EpisodeStatLines. + const TSharedPtr Root = MakeShared(); + 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() - // What about episode success? -> can be stated after the result is known - // How to invalidate data + const TSharedPtr Stats = MakeShared(); + Stats->SetObjectField("action", MakeShared(TrainingEpisodeData.ControlsStats)); + Stats->SetObjectField("observation.state", MakeShared(TrainingEpisodeData.JointsStats)); - // Anuj -> How many frames do we need to store in a single parquet chunk - // Exact data structure with correct data types - }; -} - -void UEpisodeSubSystem::SendEpisodeData(const FObservationPayload& Payload) const -{ - // PayloadBuffer.Add(Payload) - // Every X frames -> Write parquet chunk + // 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" - if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem()) - { - // Here generate the path for each image? - DataTransfer->WriteImageToDisk(CurrentRobot->PhysicsSceneProxy->GetMujocoData().time); - - // 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); - } + // Append + Root->SetObjectField("stats", Stats); + + // Serialize into FString + FString Output; + const TSharedRef> Writer = TJsonWriterFactory<>::Create(&Output); + FJsonSerializer::Serialize(Root.ToSharedRef(), Writer); + EpisodeStatLines.Add(Output); } + +void UEpisodeSubSystem::CreateEpisodeParquetFile() +{ + // 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 + + // Create a jsonl file and store in the correct directory + // concat TArray 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 +} + diff --git a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotComponent.cpp b/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotComponent.cpp index 5d796682..b5a8ddb9 100644 --- a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotComponent.cpp +++ b/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotComponent.cpp @@ -64,3 +64,19 @@ void URobotPilotComponent::SetRobotCurrentRewardZone(const FTransform& RewardTra void URobotPilotComponent::ReceiveRemoteCommand(const FRemoteControlPayload& RemoteRobotPayload) { } + +FJsonObject URobotPilotComponent::GetBufferedControlsData() +{ + return FJsonObject(); +} + +FJsonObject URobotPilotComponent::GetBufferedJointsData() +{ + return FJsonObject(); +} + +FTrainingEpisodeData URobotPilotComponent::GetTrainingEpisodeData() +{ + return FTrainingEpisodeData(); +} + diff --git a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp b/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp index 38144155..b0c56686 100644 --- a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp +++ b/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp @@ -24,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() @@ -44,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); @@ -91,11 +88,179 @@ void URobotPilotSO100Component::ReceiveRemoteCommand(const FRemoteControlPayload { for (const auto& [ActuatorName, ActuatorValue] : RemoteRobotPayload.Commands) { - // Will print an error message if it doesn't exists + // 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> ControlsArray; + + for (const auto& Control : ControlsDataBuffer) + { + TSharedPtr ControlJson = MakeShared(); + + 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(ControlJson)); + } + + BufferedData.SetArrayField(TEXT("action"), ControlsArray); + + return BufferedData; +} + +FJsonObject URobotPilotSO100Component::GetBufferedJointsData() +{ + auto BufferedData = FJsonObject(); + TArray> ControlsArray; + + for (const auto& Joint : JointsDataBuffer) + { + TSharedPtr ControlJson = MakeShared(); + + 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(ControlJson)); + } + + BufferedData.SetArrayField(TEXT("observation.state"), ControlsArray); + + return BufferedData; +} + +FJsonObject URobotPilotSO100Component::GetStats(const TArray& 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> + auto ActuatorsToJsonArray = [](const FSo100Actuators& Actuators) -> TArray> + { + return { + MakeShared(Actuators.Rotation), + MakeShared(Actuators.Pitch), + MakeShared(Actuators.Elbow), + MakeShared(Actuators.WristPitch), + MakeShared(Actuators.WristRoll), + MakeShared(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(ActuatorStates.Num()) }); + + return StatsJson; +} + void URobotPilotSO100Component::PrintCurrentActuators() const { PrintActuators(GetCurrentControlsFromPhysicScene()); diff --git a/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h b/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h index 755de0b2..6f232aac 100644 --- a/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h +++ b/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h @@ -1,16 +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 Tasks; + + UPROPERTY() + int32 Length = -1; +}; + + + + UCLASS() class LUCKYWORLDV2_API UEpisodeSubSystem : public UWorldSubsystem { @@ -30,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; @@ -51,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: @@ -60,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(); @@ -92,23 +115,26 @@ private: void InitCameras(); TArray> 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 EpisodeStatLines; + void CreateEpisodesStatsJsonFile(); + void CreateEpisodesJsonFile(); + void CreateInfoJsonFile(); + void CreateTasksJsonFile(); }; diff --git a/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotComponent.h b/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotComponent.h index b1b82568..f2108a4e 100644 --- a/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotComponent.h +++ b/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotComponent.h @@ -12,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) @@ -38,7 +49,13 @@ public: 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 diff --git a/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h b/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h index 768f7e47..aa3daec1 100644 --- a/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h +++ b/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h @@ -46,6 +46,15 @@ public: 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 ControlsDataBuffer; + TArray JointsDataBuffer; + static FJsonObject GetStats(const TArray& ActuatorStates); private: FTransform TargetTransform;