You've already forked LuckyWorld
142 lines
3.5 KiB
C++
142 lines
3.5 KiB
C++
#pragma once
|
|
#include "CoreMinimal.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
|
|
{
|
|
GENERATED_BODY()
|
|
|
|
public:
|
|
// Setup
|
|
UEpisodeSubSystem();
|
|
virtual void Initialize(FSubsystemCollectionBase& Collection) override;
|
|
virtual void Deinitialize() override;
|
|
|
|
// ----------------
|
|
// ----- TICK -----
|
|
// ----------------
|
|
// TODO I don't like this solution, it's hacky - Tick should be in a component, a primitive or scene
|
|
// TODO + it's leaking, not properly teared down
|
|
// 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;
|
|
|
|
|
|
// --------------------
|
|
// ------- DEBUG ------
|
|
// --------------------
|
|
UPROPERTY()
|
|
ATextRenderActor* DebugTextActor = nullptr;
|
|
int32 SuccessEpisodes = 0;
|
|
int32 FailEpisodes = 0;
|
|
void UpdateDebugTextActor() const;
|
|
|
|
|
|
// ---------------------
|
|
// ------- START -------
|
|
// ---------------------
|
|
/**
|
|
* Called by the UI when pressing the "Capture" button
|
|
*/
|
|
UFUNCTION(BlueprintCallable)
|
|
void StartTraining(int32 EpisodesCountIn, FString BaseImageDataPathIn, FString TaskDescriptionIn);
|
|
void EndTraining();
|
|
|
|
UPROPERTY()
|
|
ULuckyDataTransferSubsystem* DataTransfer = nullptr;
|
|
|
|
|
|
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();
|
|
|
|
// Where the robot has to place the object
|
|
FTransform EpisodeObjectBaseTransform = FTransform::Identity;
|
|
|
|
// The object that will serve for the episode
|
|
TObjectPtr<AMujocoStaticMeshActor> EpisodeTargetObject;
|
|
|
|
// ---------------------
|
|
// ------- ROBOT -------
|
|
// ---------------------
|
|
|
|
// The state of capture - if true we should call the scene capture and data transfer
|
|
bool bIsCapturing = false;
|
|
bool bIsEpisodeRunning = false;
|
|
int32 EpisodesToCapture = 0;
|
|
int32 CapturedEpisodes = 0;
|
|
|
|
void FindEpisodeObjectFromScene();
|
|
void FindRobotPawnFromScene();
|
|
|
|
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();
|
|
|
|
// 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();
|
|
};
|
|
|
|
|
|
|