FT - Looping Episodes

This commit is contained in:
Jb win
2025-05-04 01:04:44 +07:00
parent 45490051cd
commit 9806b4f5af
17 changed files with 454 additions and 132 deletions

View File

@ -4,11 +4,13 @@
#include "Kismet/GameplayStatics.h"
#include "Robot/RobotPawn.h"
#include "Robot/PilotComponent/RobotPilotComponent.h"
#include "Robot/PilotComponent/RobotPilotSO100Component.h"
#include "LuckyDataTransferSubsystem.h"
#include "Components/TextRenderComponent.h"
#include "Engine/TextRenderActor.h"
UEpisodeSubSystem::UEpisodeSubSystem()
{
}
void UEpisodeSubSystem::Initialize(FSubsystemCollectionBase& Collection)
@ -18,70 +20,148 @@ void UEpisodeSubSystem::Initialize(FSubsystemCollectionBase& Collection)
void UEpisodeSubSystem::Deinitialize()
{
bTickEnabled = false;
FTSTicker::GetCoreTicker().RemoveTicker(TickHandle);
Super::Deinitialize();
}
void UEpisodeSubSystem::Tick(float DeltaTime)
{
// if capture has started
if (!bIsCapturing || CapturedEpisodes >= EpisodesCount) return;
// 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
// Noah
// Configure the DataTransfer -> Use CurrentRobot->Cameras
// Start the Capture
// Make specs for JB to add API on the robot data
if (!bIsEpisodeRunning)
{
StartEpisode();
}
else
{
const bool bIsEpisodeCompleted = CheckEpisodeCompletion();
// JB
// Check status of the episode
// object near base location ?
// Get Robot right zone
// Get Robot left zone
// Move Object at Transform
// Robot->Grab object at Location + Deliver at Zone left/right
// Start Episode new episode - increase counter
// Check the distance of Object to BaseLocation
// EpisodeTargetObject->MainActorBody->GetComponentTransform();
// Order the robot to go fetch an object
// Both of us once we finished our own tasks and synced
// Add the timestamp debug component on the scene to check if the rendered frame and the data are in sync
if (bIsEpisodeCompleted && CapturedEpisodes <= EpisodesToCapture)
{
return StartEpisode();
}
// Here shouldn't we rewrite the frames to know if the episode was a success or a failure?
// Tick
// Get object location + compute velocity
// if velocity == 0 -> set Time since velocity 0
// if velocity == 0 for time > trigger restart time
// check object location - is in drop zone?
// if in drop zone -> set success
// Check ObjectLocation
// If Velocity
// ProceduralSceneController - is object spawned? -> If not spawn it
// if object spawned && !robot.hasTarget -> Robot -> SetTarget
// if object.IsSpawned && robot.hasTarget -> Capture and send data
// ProceduralSceneController -> is object collected?
// How to reset the episode?
// Maybe this should not be done in the tick but after episode completion
const auto Payload = CreatePayload();
SendEpisodeData(Payload);
}
}
void UEpisodeSubSystem::StartNewEpisodesSeries(const int32 EpisodesCountIn)
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();
EpisodesCount = EpisodesCountIn;
bIsCapturing = true;
CurrentObjectBaseLocation = CurrentRobot->RobotPilotComponent->GetReachableTransform();
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()
@ -106,22 +186,51 @@ void UEpisodeSubSystem::FindRobotPawnFromScene()
void UEpisodeSubSystem::ConfigureDataCapture()
{
// Noah's space of artistic expression
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();
}
}
void UEpisodeSubSystem::CreatePayload()
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();
// const auto TimeStamp = CurrentRobot->PhysicsSceneProxy->GetMujocoData().time;
// const auto So100PilotCmp = Cast<URobotPilotSO100Component>(CurrentRobot->RobotPilotComponent);
// const auto Joints = So100PilotCmp->GetCurrentControlsFromPhysicScene();
// JB
// Here I need this specific data
// 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?
// How to invalidate data
};
}
void UEpisodeSubSystem::SendEpisodeData()
void UEpisodeSubSystem::SendEpisodeData(const FObservationPayload& Payload) const
{
// Send the Data
if (ULuckyDataTransferSubsystem* DataTransfer = GetWorld()->GetSubsystem<ULuckyDataTransferSubsystem>())
{
// Here generate the path for each image?
// DataTransfer->WriteImageToDisk(BaseImageDataPath, 0.f);
// 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);
}
}