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"

UEpisodeSubSystem::UEpisodeSubSystem()
{
	
}

void UEpisodeSubSystem::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);
}

void UEpisodeSubSystem::Deinitialize()
{
	bTickEnabled = false;
	FTSTicker::GetCoreTicker().RemoveTicker(TickHandle);
	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();

		if (bIsEpisodeCompleted && CapturedEpisodes <= EpisodesToCapture)
		{
			return StartEpisode();
		}
		
		// 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);
	}
}

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

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::StartNewEpisodesSeries(const int32 EpisodesCountIn, FString BaseImageDataPathIn)
{
	// 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;

	StartTicking();
}

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();
}

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 (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
	{
		//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;
		}
	}
}

FObservationPayload UEpisodeSubSystem::CreatePayload()
{
	// 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();

	// 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? -> can be stated after the result is known
		// How to invalidate data

		// 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
	
	if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
	{
		// 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);
	}
}