2025-05-02 23:23:25 +07:00
|
|
|
|
#pragma once
|
|
|
|
|
#include "CoreMinimal.h"
|
2025-05-04 01:04:44 +07:00
|
|
|
|
#include "ObservationData.h"
|
2025-05-02 23:23:25 +07:00
|
|
|
|
#include "Subsystems/WorldSubsystem.h"
|
2025-05-04 01:04:44 +07:00
|
|
|
|
#include "Stats/Stats.h"
|
2025-05-02 23:23:25 +07:00
|
|
|
|
#include "EpisodeSubSystem.generated.h"
|
|
|
|
|
|
|
|
|
|
|
2025-05-04 01:04:44 +07:00
|
|
|
|
class ATextRenderActor;
|
2025-05-03 01:45:06 +07:00
|
|
|
|
class AMujocoStaticMeshActor;
|
2025-05-02 23:23:25 +07:00
|
|
|
|
class ARobotPawn;
|
|
|
|
|
|
|
|
|
|
UCLASS()
|
|
|
|
|
class LUCKYWORLDV2_API UEpisodeSubSystem : public UWorldSubsystem
|
|
|
|
|
{
|
|
|
|
|
GENERATED_BODY()
|
|
|
|
|
|
|
|
|
|
public:
|
|
|
|
|
// Setup
|
|
|
|
|
UEpisodeSubSystem();
|
2025-05-04 01:04:44 +07:00
|
|
|
|
virtual void Initialize(FSubsystemCollectionBase& Collection) override;
|
|
|
|
|
virtual void Deinitialize() override;
|
2025-05-02 23:23:25 +07:00
|
|
|
|
|
2025-05-04 01:04:44 +07:00
|
|
|
|
// ----------------
|
|
|
|
|
// ----- 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();
|
|
|
|
|
FTSTicker::FDelegateHandle TickHandle;
|
|
|
|
|
bool bTickEnabled = true;
|
2025-05-02 23:23:25 +07:00
|
|
|
|
|
2025-05-04 01:04:44 +07:00
|
|
|
|
|
|
|
|
|
// --------------------
|
|
|
|
|
// ------- DEBUG ------
|
|
|
|
|
// --------------------
|
|
|
|
|
UPROPERTY()
|
|
|
|
|
ATextRenderActor* DebugTextActor = nullptr;
|
|
|
|
|
int32 SuccessEpisodes = 0;
|
|
|
|
|
int32 FailEpisodes = 0;
|
|
|
|
|
void UpdateDebugTextActor() const;
|
|
|
|
|
|
|
|
|
|
|
2025-05-02 23:23:25 +07:00
|
|
|
|
// ---------------------
|
|
|
|
|
// ------- START -------
|
|
|
|
|
// ---------------------
|
|
|
|
|
/**
|
|
|
|
|
* Called by the UI when pressing the "Capture" button
|
|
|
|
|
*/
|
2025-05-04 01:04:44 +07:00
|
|
|
|
UFUNCTION(BlueprintCallable)
|
|
|
|
|
void StartNewEpisodesSeries(int32 EpisodesCountIn, FString BaseImageDataPathIn);
|
2025-05-03 01:45:06 +07:00
|
|
|
|
|
|
|
|
|
|
2025-05-02 23:23:25 +07:00
|
|
|
|
|
|
|
|
|
private:
|
2025-05-03 01:45:06 +07:00
|
|
|
|
// ---------------------
|
|
|
|
|
// ------- FLOW --------
|
|
|
|
|
// ---------------------
|
|
|
|
|
void StartEpisode();
|
2025-05-04 01:04:44 +07:00
|
|
|
|
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();
|
2025-05-03 01:45:06 +07:00
|
|
|
|
|
2025-05-04 01:04:44 +07:00
|
|
|
|
// Where the robot has to place the object
|
|
|
|
|
FTransform EpisodeObjectBaseTransform = FTransform::Identity;
|
|
|
|
|
|
|
|
|
|
// The object that will serve for the episode
|
|
|
|
|
TObjectPtr<AMujocoStaticMeshActor> EpisodeTargetObject;
|
2025-05-03 01:45:06 +07:00
|
|
|
|
|
|
|
|
|
// ---------------------
|
|
|
|
|
// ------- ROBOT -------
|
|
|
|
|
// ---------------------
|
|
|
|
|
|
|
|
|
|
// The state of capture - if true we should call the scene capture and data transfer
|
2025-05-02 23:23:25 +07:00
|
|
|
|
bool bIsCapturing = false;
|
2025-05-04 01:04:44 +07:00
|
|
|
|
bool bIsEpisodeRunning = false;
|
|
|
|
|
int32 EpisodesToCapture = 0;
|
2025-05-03 01:45:06 +07:00
|
|
|
|
int32 CapturedEpisodes = 0;
|
|
|
|
|
|
|
|
|
|
void FindEpisodeObjectFromScene();
|
|
|
|
|
void FindRobotPawnFromScene();
|
2025-05-02 23:23:25 +07:00
|
|
|
|
|
|
|
|
|
UPROPERTY()
|
|
|
|
|
TObjectPtr<ARobotPawn> CurrentRobot;
|
2025-05-03 01:45:06 +07:00
|
|
|
|
|
2025-05-04 01:04:44 +07:00
|
|
|
|
|
2025-05-03 01:45:06 +07:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
// --------------------
|
|
|
|
|
// ------- DATA -------
|
|
|
|
|
// --------------------
|
2025-05-04 01:04:44 +07:00
|
|
|
|
FString BaseImageDataPath;
|
|
|
|
|
|
2025-05-03 01:45:06 +07:00
|
|
|
|
// Noah here add anything you need
|
|
|
|
|
void ConfigureDataCapture();
|
|
|
|
|
|
2025-05-04 01:04:44 +07:00
|
|
|
|
FObservationPayload CreatePayload();
|
2025-05-03 01:45:06 +07:00
|
|
|
|
|
2025-05-04 01:04:44 +07:00
|
|
|
|
void SendEpisodeData(const FObservationPayload& Payload) const;
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2025-05-02 23:23:25 +07:00
|
|
|
|
};
|
2025-05-03 01:45:06 +07:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|