LuckyWorld/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp

#include "Episode/EpisodeSubSystem.h"
#include "Actors/MujocoStaticMeshActor.h"
#include "Actors/MujocoVolumeActor.h"
#include "Kismet/GameplayStatics.h"
#include "Robot/RobotPawn.h"
#include "Robot/PilotComponent/RobotPilotComponent.h"
#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()
{
	
}

void UEpisodeSubSystem::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);
	if (ULuckyDataTransferSubsystem* DataTransferSubSystem = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
	{
		DataTransfer = DataTransferSubSystem;
	}
}

void UEpisodeSubSystem::Deinitialize()
{
	StopTicking();
	Super::Deinitialize();
}

void UEpisodeSubSystem::Tick(float DeltaTime)
{
	// TODO we want to get this outside of the Tick
	if (!bTickEnabled) return;
	
	// If no robot or no object
	if (!EpisodeTargetObject || !CurrentRobot) return;
	
	// if capture hasn't started
	if (!bIsCapturing || CapturedEpisodes >= EpisodesToCapture) return;

	// Here we are capturing the data, running an episode
	if (!bIsEpisodeRunning)
	{
		StartEpisode();
	}
	else
	{
		const bool bIsEpisodeCompleted = CheckEpisodeCompletion();

		// Write Image on the disk
		if (DataTransfer) DataTransfer->WriteImageToDisk(CurrentRobot->PhysicsSceneProxy->GetMujocoData().time);
		EpisodeFrames++;

		if (bIsEpisodeCompleted && CapturedEpisodes <= EpisodesToCapture)
		{
			EndEpisode();
			StartEpisode();
		}
		else
		{
			EndTraining();
		}
	}
}

void UEpisodeSubSystem::StartTicking()
{
	const FTickerDelegate TickDelegate = FTickerDelegate::CreateLambda([this](const float DeltaTime)
	{
		Tick(DeltaTime);
		return bTickEnabled;
	});
	TickHandle = FTSTicker::GetCoreTicker().AddTicker(TickDelegate);
}

void UEpisodeSubSystem::StopTicking()
{
	bTickEnabled = false;
	FTSTicker::GetCoreTicker().RemoveTicker(TickHandle);
}

void UEpisodeSubSystem::UpdateDebugTextActor() const
{
	if (!IsValid(DebugTextActor)) return;
	
	const auto TextRender = DebugTextActor->GetTextRender();
	const FString Txt = FString::Printf(TEXT("Episodes run: %i \nSuccess: %i \nFailed: %i"), CapturedEpisodes, SuccessEpisodes, FailEpisodes);
	TextRender->SetText(FText::FromString(Txt));
}

void UEpisodeSubSystem::StartTraining(const int32 EpisodesCountIn, FString BaseImageDataPathIn, FString TaskDescriptionIn)
{
	// Debug
	const auto DebugTextActorPtr = UGameplayStatics::GetActorOfClass(this->GetWorld(), ATextRenderActor::StaticClass());
	if (DebugTextActorPtr && Cast<ATextRenderActor>(DebugTextActorPtr))
	{
		DebugTextActor = Cast<ATextRenderActor>(DebugTextActorPtr);
	}
	
	// Robot and Exercise
	FindEpisodeObjectFromScene();
	FindRobotPawnFromScene();
	EpisodesToCapture = EpisodesCountIn;
	SuccessEpisodes = 0;
	FailEpisodes = 0;
	StartEpisode();

	// 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 
	if (!CurrentRobot->RobotPilotComponent->GetIsReadyForTraining()) return;

	// Let's hardcode this for now, and figure out later how to do it correctly with Anuj/Ethan inputs
	const FTransform RobotTransform = CurrentRobot->RobotActor->GetActorTransform();
	constexpr float HardCodedRewardDistanceFromRobotPivot = 15.f; // TODO This should not be hardcoded as it depends from robot type
	EpisodeRewardZone = FTransform{
		// TODO RobotArm right is the forward vector due to rotation the Robot -90 yaw at robot spawn - FIX ME 
		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();
	
	// Move Scenario Object to its location - Done in the PhysicsScene
	CurrentRobot->PhysicsSceneProxy->UpdateGeomTransform(EpisodeTargetObject->MainActorBody.GetName(), EpisodeObjectBaseTransform);
	
	// Set Target on the bot - it will go grab the object
	CurrentRobot->RobotPilotComponent->SetRobotTarget(EpisodeObjectBaseTransform);
	CurrentRobot->RobotPilotComponent->SetRobotCurrentRewardZone(EpisodeRewardZone);
	
	// Enable Tick checks
	bIsEpisodeRunning = true;
	bIsCapturing = true;

	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());
	const auto Loc = GeomTransform.GetLocation();

	const auto DistanceFromStart = FVector::Distance(EpisodeObjectBaseTransform.GetLocation(), Loc);

	if (DistanceFromStart <= 2) return false; // Episode is running

	// TODO This can be used to early detect episode failure and restart the episode faster 
	const auto DotUp = FVector::DotProduct(FVector::UpVector, GeomTransform.GetRotation().GetUpVector());  
	
	// Robot did not finish the episode yet 
	if (!CurrentRobot->RobotPilotComponent->GetIsInRestState()) return false;
	
	// Here we are away from Start zone and Robot has finished the exercise
	const auto DistanceToReward = FVector::Distance(EpisodeRewardZone.GetLocation(), Loc);
	if (DistanceToReward < EpisodeRewardZoneRadius)
	{
		SuccessEpisodes++;
	}
	else
	{
		FailEpisodes++;	
	}

	CapturedEpisodes++;
	return true;
}

void UEpisodeSubSystem::FindEpisodeObjectFromScene()
{
	TArray<AActor*> MujocoObjects;
	UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), AMujocoStaticMeshActor::StaticClass(), MujocoObjects);
	if (MujocoObjects.IsValidIndex(0) && Cast<AMujocoStaticMeshActor>(MujocoObjects[0]))
	{
		EpisodeTargetObject = Cast<AMujocoStaticMeshActor>(MujocoObjects[0]);		
	}
}

void UEpisodeSubSystem::FindRobotPawnFromScene()
{
	TArray<AActor*> RobotPawns;
	UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), ARobotPawn::StaticClass(), RobotPawns);
	if (RobotPawns.IsValidIndex(0) && Cast<ARobotPawn>(RobotPawns[0]))
	{
		CurrentRobot = Cast<ARobotPawn>(RobotPawns[0]);		
	}
}

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 (!DataTransfer) return;
	DataTransfer->CreateCaptureSessionID();
	InitCameras();
	for (const auto& Cam : Cameras)
	{
		DataTransfer->RegisterSensor(Cam.Get());
		Cam->SensorInfo.bActive = true;
	}
}

void UEpisodeSubSystem::CreateEpisodeStatJsonLine(const FTrainingEpisodeData& TrainingEpisodeData)
{
	// EpisodeStatLines.
	const TSharedPtr<FJsonObject> Root = MakeShared<FJsonObject>();
	Root->SetNumberField("episode_index", CapturedEpisodes);

	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<>> 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<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
}