FT - Looping Episodes

Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Actors/MujocoStaticMeshActor.cpp

@@ -1,9 +1,7 @@
 #include "Actors/MujocoStaticMeshActor.h"
 #include "EngineUtils.h"
 #include "Components/MujocoBodyComponent.h"
-#include "Components/MujocoActuatorComponent.h"
-#include "Components/MujocoEqualityComponent.h"
-#include "Components/MujocoTendonComponent.h"
+#include "Components/MujocoGeomComponent.h"
 
 AMujocoStaticMeshActor::AMujocoStaticMeshActor()
 {
@@ -14,8 +12,22 @@ void AMujocoStaticMeshActor::BeginPlay()
 {
 	Super::BeginPlay();
 	MainActorBody = Cast<UMujocoBodyComponent>(this->GetComponentByClass(UMujocoBodyComponent::StaticClass()));
+	MainActorBodyGeometry = Cast<UMujocoGeomComponent>(this->GetComponentByClass(UMujocoGeomComponent::StaticClass()));	
 }
 
+FTransform AMujocoStaticMeshActor::GetGeometryTransform()
+{
+	if (!MainActorBodyGeometry) return FTransform::Identity;
+	
+	return MainActorBodyGeometry->GetComponentTransform();
+}
+
+// ------------------------------
+// ------------------------------
+// -------- EDITOR ONLY ---------
+// ------------------------------
+// ------------------------------
+
 #if WITH_EDITOR
 
 void AMujocoStaticMeshActor::PostDuplicate(EDuplicateMode::Type DuplicateMode)

Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Actors/MujocoVolumeActor.cpp

@@ -209,7 +209,7 @@ mjData_& AMujocoVolumeActor::GetMujocoData() const
 	return *MujocoData.Get();
 }
 
-void AMujocoVolumeActor::UpdateGeomPosition(const FString BodyName, const FVector& NewPosition, const FQuat& NewRotation)
+void AMujocoVolumeActor::UpdateGeomTransform(const FString& BodyName, const FTransform& NewTransform)
 {
 	// Step 1: Get body ID
 	const int Body_ID = mj_name2id(MujocoModel.Get(), mjOBJ_BODY, TCHAR_TO_ANSI(*BodyName));
@@ -230,15 +230,15 @@ void AMujocoVolumeActor::UpdateGeomPosition(const FString BodyName, const FVecto
 	const int Qvel_Adr = MujocoModel->jnt_dofadr[Joint_Adr];
 
 	// Step 4: Convert position and rotation
-	MujocoData->qpos[Qpos_Adr + 0] = NewPosition.X / 100.f;       // X
-	MujocoData->qpos[Qpos_Adr + 1] = -NewPosition.Y / 100.f;      // Y (flip for Unreal Z-up)
-	MujocoData->qpos[Qpos_Adr + 2] = NewPosition.Z / 100.f;       // Z
+	MujocoData->qpos[Qpos_Adr + 0] = NewTransform.GetLocation().X / 100.f;       // X
+	MujocoData->qpos[Qpos_Adr + 1] = -NewTransform.GetLocation().Y / 100.f;      // Y (flip for Unreal Z-up)
+	MujocoData->qpos[Qpos_Adr + 2] = NewTransform.GetLocation().Z / 100.f;       // Z
 
 	// Unreal (X, Y, Z, W) → MuJoCo (W, X, Y, Z)
-	MujocoData->qpos[Qpos_Adr + 3] = NewRotation.W;
-	MujocoData->qpos[Qpos_Adr + 4] = NewRotation.X;
-	MujocoData->qpos[Qpos_Adr + 5] = NewRotation.Y;
-	MujocoData->qpos[Qpos_Adr + 6] = NewRotation.Z;
+	MujocoData->qpos[Qpos_Adr + 3] = NewTransform.GetRotation().W;
+	MujocoData->qpos[Qpos_Adr + 4] = NewTransform.GetRotation().X;
+	MujocoData->qpos[Qpos_Adr + 5] = -NewTransform.GetRotation().Y;
+	MujocoData->qpos[Qpos_Adr + 6] = -NewTransform.GetRotation().Z;
 
 	// Step 5: Zero velocity
 	for (int i = 0; i < 6; i++) {
@@ -249,6 +249,43 @@ void AMujocoVolumeActor::UpdateGeomPosition(const FString BodyName, const FVecto
 	mj_forward(MujocoModel.Get(), MujocoData.Get());
 }
 
+FTransform AMujocoVolumeActor::GetGeometryTransform(const FString& BodyName) const
+{
+	if (!MujocoModel.IsValid() || !MujocoData.IsValid()) return FTransform::Identity;
+
+	// Step 1: Get body ID
+	const int Body_ID = mj_name2id(MujocoModel.Get(), mjOBJ_BODY, TCHAR_TO_ANSI(*BodyName));
+	if (Body_ID < 0) {
+		UE_LOG(LogTemp, Error, TEXT("Body not found: %s"), *BodyName);
+		return FTransform::Identity;
+	}
+
+	// Step 2: Get the joint ID
+	const int Joint_Adr = MujocoModel->body_jntadr[Body_ID];
+	if (MujocoModel->jnt_type[Joint_Adr] != mjJNT_FREE) {
+		UE_LOG(LogTemp, Error, TEXT("Body '%s' does not have a free joint."), *BodyName);
+		return FTransform::Identity;
+	}
+
+	// Step 3: Get qpos address
+	const int Qpos_Adr = MujocoModel->jnt_qposadr[Joint_Adr];
+
+	// Step 4: Extract position (convert from MuJoCo to Unreal scale and axis)
+	FVector Position;
+	Position.X = MujocoData->qpos[Qpos_Adr + 0] * 100.f;        // X
+	Position.Y = -MujocoData->qpos[Qpos_Adr + 1] * 100.f;       // Y (flip back)
+	Position.Z = MujocoData->qpos[Qpos_Adr + 2] * 100.f;        // Z
+
+	// Step 5: Extract rotation (MuJoCo (W, X, Y, Z) → Unreal (X, Y, Z, W))
+	FQuat Rotation;
+	Rotation.W = MujocoData->qpos[Qpos_Adr + 3];
+	Rotation.X = MujocoData->qpos[Qpos_Adr + 4];
+	Rotation.Y = -MujocoData->qpos[Qpos_Adr + 5];
+	Rotation.Z = -MujocoData->qpos[Qpos_Adr + 6];
+
+	return FTransform(Rotation, Position);
+}
+
 void AMujocoVolumeActor::SetActuatorValue(const FString& ActuatorName, double Value)
 {
 	if (MujocoModel)

Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoStaticMeshActor.h

@@ -1,10 +1,11 @@
 #pragma once
 
 #include "CoreMinimal.h"
-#include "Engine/StaticMeshActor.h"
 #include "MujocoStaticMeshActor.generated.h"
 
+class UMujocoGeomComponent;
 class UMujocoBodyComponent;
+
 /**
  *
  */
@@ -18,11 +19,22 @@ public:
 
 	virtual void BeginPlay() override;
 	
+	TObjectPtr<UMujocoBodyComponent> MainActorBody;
+	FTransform GetGeometryTransform();
+	
+private:
+	TObjectPtr<UMujocoGeomComponent> MainActorBodyGeometry;
+	
+	// ------------------------------
+	// ------------------------------
+	// -------- EDITOR ONLY ---------
+	// ------------------------------
+	// ------------------------------
 #if WITH_EDITOR
+public:
 	virtual void PreDuplicate(FObjectDuplicationParameters& DupParams) override;
 
 	virtual void PostDuplicate(EDuplicateMode::Type DuplicateMode) override;
 #endif
 
-	TObjectPtr<UMujocoBodyComponent> MainActorBody;
 };

Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoVolumeActor.h

@@ -111,7 +111,9 @@ public:
 	mjData_& GetMujocoData() const;
 
 	UFUNCTION(BlueprintCallable, Category = "Mujoco")
-	void UpdateGeomPosition(const FString BodyName, const FVector& NewPosition, const FQuat& NewRotation);
+	void UpdateGeomTransform(const FString& BodyName, const FTransform& NewTransform);
+
+	FTransform GetGeometryTransform(const FString& BodyName) const;
 	
 	// ---------------------------
 	// ------- POST UPDATE -------

Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp

@@ -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
+	if (!bIsEpisodeRunning)
+	{
+		StartEpisode();
+	}
+	else
+	{
+		const bool bIsEpisodeCompleted = CheckEpisodeCompletion();
 
-	// Noah
-	// Configure the DataTransfer -> Use CurrentRobot->Cameras
-	// Start the Capture
-	// Make specs for JB to add API on the robot data
-
-	// 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
-
-	// 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?
+		if (bIsEpisodeCompleted && CapturedEpisodes <= EpisodesToCapture)
+		{
+			return StartEpisode();
 		}
 		
-void UEpisodeSubSystem::StartNewEpisodesSeries(const int32 EpisodesCountIn)
+		// 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();
-	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
 {
+	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);
+	}
 }

Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotComponent.cpp

@@ -40,3 +40,23 @@ FTransform URobotPilotComponent::GetReachableTransform()
 {
 	return FTransform::Identity;
 }
+
+bool URobotPilotComponent::GetIsReadyForTraining()
+{
+	// Overriden in individual components
+	return false;
+}
+
+bool URobotPilotComponent::GetIsInRestState()
+{
+	// Overriden in individual components
+	return true;
+}
+
+void URobotPilotComponent::SetRobotTarget(const FTransform& TargetTransformIn)
+{
+}
+
+void URobotPilotComponent::SetRobotCurrentRewardZone(const FTransform& RewardTransformIn)
+{
+}

Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp

@@ -12,6 +12,11 @@ URobotPilotSO100Component::URobotPilotSO100Component()
 void URobotPilotSO100Component::BeginPlay()
 {
 	Super::BeginPlay();
+
+	// Start the RestPost animation 0.1s after loading
+	AnimationStartTime = .1f; 
+	bBreakAfterAnimation = true;
+	RestPose();
 }
 
 void URobotPilotSO100Component::TickComponent(float DeltaTime, enum ELevelTick TickType,
@@ -25,23 +30,65 @@ void URobotPilotSO100Component::TickComponent(float DeltaTime, enum ELevelTick T
 
 FTransform URobotPilotSO100Component::GetReachableTransform()
 {
-	const auto RobotTransform = RobotOwner->GetActorTransform();
-	const float Yaw = MaxYaw * FMath::RandRange(0., 1.) * (FMath::RandBool() ? 1 : -1);
-	const float Range = MaxRange * FMath::RandRange(0., 1.);
-	const FRotator RandomRotator = FRotator{0,0,Yaw};
-	const FVector RandomLocation = RandomRotator.RotateVector(RobotTransform.GetRotation().GetForwardVector() * Range);
-	const FRotator RandomRotation = UKismetMathLibrary::MakeRotFromXZ(RandomLocation- RobotTransform.GetLocation(), FVector::UpVector);
-	return FTransform(RandomRotation, RandomLocation);
+	const auto RobotTransform = RobotOwner->RobotActor->GetActorTransform();
+
+	// Robot actor is built Y axis forward, like everything in Unreal - Rotate to get Arm facing X world axis aka forward vector
+	const auto ArmWorldRotation = FRotator{0,90,0}.Quaternion() * RobotTransform.GetRotation();
+
+	// Find Arm Pivot Location
+	const auto ArmPivotLocation = RobotTransform.GetLocation() + RobotTransform.GetRotation().RotateVector(PivotOffset) + FVector{0,0,5};
+
+	// Compute a random Yaw
+	const float RandomYaw = MaxYaw * FMath::RandRange(0., 1.) * (FMath::RandBool() ? 1 : -1);
+	const FQuat RandomRotation = FRotator{0,RandomYaw,0}.Quaternion() * ArmWorldRotation;
+
+
+
+	// Compute Random Range within reach of the arm and add this to pivot location
+	// Add a bit more than the Jaw Offset - TODO Offsets must be better computed
+	const float RandomRange = JawOffset.X + MaxRange * FMath::RandRange(0.1f, 1.f);
+	const FVector RandomLocation = ArmPivotLocation + RandomRotation.GetForwardVector() * RandomRange;
+
+	// Find Look at rotation from location to Pivot
+	auto RewardAxis = RandomLocation-ArmPivotLocation;
+	RewardAxis.Z = 0; // Nullify Z to keep a 2D vector -> ensure the geometry roll/pitch are 0
+	const FRotator TowardPivotRotation = UKismetMathLibrary::MakeRotFromXZ(RewardAxis, FVector::UpVector);
+
+	// Debug
+	// DrawDebugLine(this->GetWorld(), ArmPivotLocation + ArmWorldRotation.GetForwardVector() * 70, ArmPivotLocation, FColor::Green, true);
+	// DrawDebugLine(this->GetWorld(), ArmPivotLocation + FVector::UpVector * 70, ArmPivotLocation, FColor::Red, true);
+	// DrawDebugLine(this->GetWorld(), RandomLocation, RandomLocation + TowardPivotRotation.Quaternion().GetForwardVector() * -50 , FColor::Blue, true);
+	
+	// Return the Object Transform
+	return FTransform(TowardPivotRotation, RandomLocation);
+}
+
+bool URobotPilotSO100Component::GetIsReadyForTraining()
+{
+	const auto CurrentJoints = GetCurrentJointsFromPhysicsScene();
+	return AreActuatorsAlmostEqual(CurrentJoints, ActuatorsRestPosition); 
+}
+
+bool URobotPilotSO100Component::GetIsInRestState()
+{
+	return CurrentAnimationState == 0 && GetIsReadyForTraining();
 }
 
 
-void URobotPilotSO100Component::SetTarget(const FTransform& TargetTransformIn)
+void URobotPilotSO100Component::SetRobotTarget(const FTransform& TargetTransformIn)
 {
 	// Set Base Values
 	TargetTransform = TargetTransformIn;
 	NextAnimationState();
 }
 
+void URobotPilotSO100Component::SetRobotCurrentRewardZone(const FTransform& RewardTransformIn)
+{
+	const auto DirectionToTarget = (RewardTransformIn.GetLocation() - RobotOwner->RobotActor->GetActorLocation()).GetSafeNormal();
+
+	// TODO This is wrong way to do because GetRightVector is the arm forward vector due to robot actor being rotated -90 yaw at spawn
+	bDropZoneIsRight = FVector::DotProduct(RobotOwner->RobotActor->GetActorRotation().Quaternion().GetRightVector(), DirectionToTarget) > 0.f;
+}
 
 
 void URobotPilotSO100Component::PrintCurrentActuators() const
@@ -110,6 +157,18 @@ float URobotPilotSO100Component::GetDeltaSumBetweenActuatorValues(const FSo100Ac
 	return DeltaSum;
 }
 
+bool URobotPilotSO100Component::AreActuatorsAlmostEqual(const FSo100Actuators& A, const FSo100Actuators& B)
+{
+	if (FMath::Abs(A.Rotation) - FMath::Abs(B.Rotation) > 0.001) return false;
+	if (FMath::Abs(A.Pitch) - FMath::Abs(B.Pitch) > 0.001) return false;
+	if (FMath::Abs(A.Elbow) - FMath::Abs(B.Elbow) > 0.001) return false;
+	if (FMath::Abs(A.WristPitch) - FMath::Abs(B.WristPitch) > 0.001) return false;
+	if (FMath::Abs(A.WristRoll) - FMath::Abs(B.WristRoll) > 0.001) return false;
+	if (FMath::Abs(A.Jaw) - FMath::Abs(B.Jaw) > 0.001) return false;
+
+	return true;
+}
+
 double URobotPilotSO100Component::GetControlJointDeltaForActuator(FString ActuatorName) const
 {
 	auto const Control = RobotOwner->PhysicsSceneProxy->GetActuatorValue(ActuatorName);
@@ -144,12 +203,25 @@ void URobotPilotSO100Component::PostPhysicStepUpdate(const float SimulationTime)
 
 		if (bHasFinishedAnimation)
 		{
+			if (bBreakAfterAnimation)
+			{
+				bBreakAfterAnimation = false;
+				AnimationStartTime = 0;
+				return;
+			}
+			
 			// AnimationStartTime = 0.; // Only for debug, can be left here but useless in normal operation mode
-			NextAnimationState();
+			return NextAnimationState();
 		}
 	}
 }
 
+bool URobotPilotSO100Component::IsJawOverCurrent() const
+{
+	const auto DeltaActJntJaw = GetControlJointDeltaForActuator(Actuator_Jaw); 
+	return DeltaActJntJaw > OverCurrentThreshold;
+}
+
 void URobotPilotSO100Component::NextAnimationState()
 {
 	const float CurrentPhysicsEngineSceneTime = RobotOwner->PhysicsSceneProxy->GetMujocoData().time;
@@ -169,15 +241,15 @@ void URobotPilotSO100Component::NextAnimationState()
 		case 4:
 			return OpenJaw();
 		default:
-			return BasePose();
+			return RestPose();
 	}
 }
 
-void URobotPilotSO100Component::BasePose()
+void URobotPilotSO100Component::RestPose()
 {
 	UE_LOG(LogTemp, Log, TEXT("Animate -> BasePose"));
 	CurrentAnimationState = 0;
-	AnimationDuration = 1.5f;
+	AnimationDuration = .5f;
 	AnimTargetRobotActuators = ActuatorsRestPosition;
 }
 
@@ -204,7 +276,14 @@ void URobotPilotSO100Component::RotateToTarget()
 	
 	// reduce/increase Yaw to not have the fixed jaw colliding with the shape - TODO use middle of the jaw instead of the wall of the jaw
 	const auto Dot = FVector::DotProduct(RotationToTarget.Quaternion().GetForwardVector(), WorldTransform.GetRotation().GetForwardVector());
-	const auto Mod = .1 * (Dot > 0 ? 1 : -1);
+
+	// TODO Better computation of Jaw Center would avoid that ugliness
+	// TODO This is not working for values Dot<0.15 - we need a better solution
+	const auto ModBaseAlpha = FMath::Abs(Dot) / 0.3;
+	const auto ModBase = 0.1 * FMath::Lerp(5.f, 1.f, FMath::Min(ModBaseAlpha, 1)); // TODO Dot Product below 0.5 needs more yaw modification 
+	
+	const auto Mod = ModBase * (Dot > 0 ? 1 : -1); // TODO Hardcoded value - compute better Jaw Offset and robot geometry awareness 
+	UE_LOG(LogTemp, Log, TEXT("Dot : %f | ModBaseAlpha: %f | ModBase: %f | Mod: %f"), Dot, ModBaseAlpha, ModBase, Mod);
 
 	// Convert to radians
 	const auto ActuatorRotation = RotationToTarget.Yaw * (1+Mod) / 180.0f * -PI; // Looks like we are not in the same referential hence the -PI instead of PI ! 
@@ -213,7 +292,7 @@ void URobotPilotSO100Component::RotateToTarget()
 	AnimTargetRobotActuators = AnimStartRobotActuators;
 	AnimTargetRobotActuators.Rotation = ActuatorRotation;
 	CurrentAnimationState = 1;
-	AnimationDuration = .7f;
+	AnimationDuration = .33f;
 }
 
 void URobotPilotSO100Component::MoveToTarget()
@@ -224,6 +303,7 @@ void URobotPilotSO100Component::MoveToTarget()
 	const auto WorldTransform = RobotOwner->RobotActor->GetActorTransform();
 	const FVector PivotWorldLocation = WorldTransform.GetLocation() + WorldTransform.GetRotation().RotateVector(PivotOffset);
 
+	// TODO Better computations
 	// Rotate Jaw offset towards target
 	// Get pure 2d rotation
 	RotationToTarget.Pitch = 0;
@@ -232,7 +312,8 @@ void URobotPilotSO100Component::MoveToTarget()
 	const auto JawPositionWorld = RotationToTarget.RotateVector(JawOffsetToPivot) + PivotWorldLocation;
 	const auto Distance = FVector::Distance(JawPositionWorld, TargetTransform.GetLocation());
 
-	const auto AlphaExtend = FMath::Clamp(Distance / MaxRange, 0., 1.);
+	// TODO Compute correct ranges and avoid to add a hardcoded value to adjust the arm extension
+	const auto AlphaExtend = FMath::Clamp(Distance / (MaxRange + 4), 0., 1.);
 	
 	// Set the target actuators values 
 	AnimTargetRobotActuators = GetCurrentJointsFromPhysicsScene();
@@ -244,15 +325,7 @@ void URobotPilotSO100Component::MoveToTarget()
 	
 	// Start the animation 	
 	CurrentAnimationState = 2;
-	AnimationDuration = 2.f;
-	
-	DrawDebugLine(
-		this->GetWorld(),
-		JawPositionWorld,
-		JawPositionWorld + (TargetTransform.GetLocation() - JawPositionWorld).GetSafeNormal() * Distance,
-		FColor::Green,
-		true
-	);
+	AnimationDuration = .66f;
 }
 
 void URobotPilotSO100Component::CloseJaw()
@@ -268,7 +341,7 @@ void URobotPilotSO100Component::CloseJaw()
 	// Start the animation
 	bDetectOverCurrent = true;
 	CurrentAnimationState = 3;
-	AnimationDuration = 2.f;
+	AnimationDuration = .5f;
 }
 
 void URobotPilotSO100Component::MoveToDropZone()
@@ -276,10 +349,13 @@ void URobotPilotSO100Component::MoveToDropZone()
 	UE_LOG(LogTemp, Log, TEXT("Animate -> MoveToDropZone"));
 
 	AnimTargetRobotActuators = ActuatorsDropZone;
-	AnimTargetRobotActuators.Rotation = ActuatorsDropZone.Rotation * (FMath::RandBool() ? 1. : -1.); 
-	AnimTargetRobotActuators.Jaw = GetCurrentJointsFromPhysicsScene().Jaw;
+	AnimTargetRobotActuators.Rotation = ActuatorsDropZone.Rotation * (bDropZoneIsRight ? 1. : -1.); 
+
+	// Here the Jaw should keep being target to closed
+	AnimTargetRobotActuators.Jaw = ClosedJaw;
+	
 	CurrentAnimationState = 4;
-	AnimationDuration = 3.f;
+	AnimationDuration = 1.5f;
 }
 
 void URobotPilotSO100Component::OpenJaw()
@@ -288,7 +364,7 @@ void URobotPilotSO100Component::OpenJaw()
 
 	AnimTargetRobotActuators.Jaw = OpenedJaw;
 	CurrentAnimationState = 5;
-	AnimationDuration = 0.6f;
+	AnimationDuration = 0.2f;
 }
 
 bool URobotPilotSO100Component::AnimateActuators(const float SimulationTime)
@@ -296,20 +372,19 @@ bool URobotPilotSO100Component::AnimateActuators(const float SimulationTime)
 	const double AnimAlpha = FMath::Clamp((SimulationTime - AnimationStartTime) / AnimationDuration, 0., 1.);
 
 	// Need to wait for the joints to be in the right position before switching to next animation
-	const auto DeltaSum = GetDeltaSumBetweenActuatorValues(AnimTargetRobotActuators, GetCurrentJointsFromPhysicsScene());
+	const bool bIsArrivedAtActuatorsConfig = AreActuatorsAlmostEqual(AnimTargetRobotActuators, GetCurrentJointsFromPhysicsScene());
 
 	// Alternative Animation completion event - checking for over-current
-	if (bDetectOverCurrent && GetControlJointDeltaForActuator(Actuator_Jaw) > OverCurrentThreshold)
+	if (bDetectOverCurrent && IsJawOverCurrent())
 	{
 		bDetectOverCurrent = false;
 		return true;
 	}
 
-	
-	// UE_LOG(LogTemp, Log, TEXT("AnimationAlpha: %f - Delta: %f"), AnimAlpha, DeltaSum);
+	// UE_LOG(LogTemp, Log, TEXT("Animate -> AnimateActuators %f - %f"), AnimAlpha, GetControlJointDeltaForActuator(Actuator_Jaw));
 	
 	// Stop the animation if we reached the target
-	if (AnimAlpha >= 1. && DeltaSum <= .001) return true;
+	if (AnimAlpha >= 1. && bIsArrivedAtActuatorsConfig) return true;
 
 	// Rotation
 	RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_Rotation, FMath::Lerp(
@@ -347,11 +422,16 @@ bool URobotPilotSO100Component::AnimateActuators(const float SimulationTime)
 	));
 
 	// Jaw
+	// If the target Jaw is closed position, aka we are grabbing, but we are over-current, then we don't hold more to not squeeze the object
+	// But if the over-current stops, aka the object rotated a bit, then we hold tighter
+	if (AnimTargetRobotActuators.Jaw != ClosedJaw || !IsJawOverCurrent())
+	{
 		RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_Jaw, FMath::Lerp(
 			AnimStartRobotActuators.Jaw,
 			AnimTargetRobotActuators.Jaw,
 			AnimAlpha
 		));
+	}
 
 	return false;
 }

Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp

@@ -14,7 +14,7 @@ ARobotPawn::ARobotPawn()
 void ARobotPawn::BeginPlay()
 {
 	Super::BeginPlay();
-	// InitRobot(); // TODO Maybe move to GameInstance to control when we initialize the robot completely
+	InitRobot(); // TODO Maybe move to GameInstance to control when we initialize the robot completely
 }
 
 void ARobotPawn::InitRobot()

Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h

@@ -1,9 +1,12 @@
 #pragma once
 #include "CoreMinimal.h"
+#include "ObservationData.h"
 #include "Subsystems/WorldSubsystem.h"
+#include "Stats/Stats.h"
 #include "EpisodeSubSystem.generated.h"
 
 
+class ATextRenderActor;
 class AMujocoStaticMeshActor;
 class ARobotPawn;
 
@@ -15,10 +18,30 @@ class LUCKYWORLDV2_API UEpisodeSubSystem : public UWorldSubsystem
 public:
 	// Setup
 	UEpisodeSubSystem();
-	virtual void Initialize(FSubsystemCollectionBase& Collection);
-	virtual void Deinitialize();
+	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();
+	FTSTicker::FDelegateHandle TickHandle;
+	bool bTickEnabled = true;
+
+
+	// --------------------
+	// ------- DEBUG ------
+	// --------------------
+	UPROPERTY()
+	ATextRenderActor* DebugTextActor = nullptr;
+	int32 SuccessEpisodes = 0;
+	int32 FailEpisodes = 0;
+	void UpdateDebugTextActor() const;
 	
-	virtual void Tick(float DeltaTime);
 	
 	// ---------------------
 	// ------- START -------
@@ -26,7 +49,8 @@ public:
 	/**
 	 * Called by the UI when pressing the "Capture" button
 	 */
-	void StartNewEpisodesSeries(int32 EpisodesCountIn);
+	UFUNCTION(BlueprintCallable)
+	void StartNewEpisodesSeries(int32 EpisodesCountIn, FString BaseImageDataPathIn);
 
 	
 
@@ -35,7 +59,15 @@ private:
 	// ------- FLOW --------
 	// ---------------------
 	void StartEpisode();
+	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 -------
@@ -43,7 +75,8 @@ private:
 	
 	// The state of capture - if true we should call the scene capture and data transfer
 	bool bIsCapturing = false;
-	int32 EpisodesCount = 0;
+	bool bIsEpisodeRunning = false;
+	int32 EpisodesToCapture = 0;
 	int32 CapturedEpisodes = 0;
 	
 	void FindEpisodeObjectFromScene();
@@ -52,22 +85,24 @@ private:
 	UPROPERTY()
 	TObjectPtr<ARobotPawn> CurrentRobot;
 
-	// The object that will serve for the episode
-	TObjectPtr<AMujocoStaticMeshActor> EpisodeTargetObject;
-	FTransform CurrentObjectBaseLocation = FTransform::Identity;
-	FTransform CurrentObjectTargetLocation = FTransform::Identity;
+
 
 
 	// --------------------
 	// ------- DATA -------
 	// --------------------
+	FString BaseImageDataPath;
 	
 	// Noah here add anything you need
 	void ConfigureDataCapture();
 
-	void CreatePayload();
+	FObservationPayload CreatePayload();
+
+	void SendEpisodeData(const FObservationPayload& Payload) const;
+
+
+
 	
-	void SendEpisodeData();
 };
 
 

Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotComponent.h

@@ -27,6 +27,13 @@ public:
 	virtual void InitPilotComponent();
 	virtual void PostPhysicStepUpdate(const float SimulationTime);
 	virtual FTransform GetReachableTransform();
+	virtual bool GetIsReadyForTraining();
+	virtual bool GetIsInRestState();
+	
+	UFUNCTION(BlueprintCallable)
+	virtual void SetRobotTarget(const FTransform& TargetTransformIn);
+	virtual void SetRobotCurrentRewardZone(const FTransform& RewardTransformIn);
+		
 	
 protected: // Child class need access
 	// Only to easy access within the component

Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h

@@ -37,13 +37,16 @@ public:
 	virtual void BeginPlay() override;
 	virtual void TickComponent(float DeltaTime, enum ELevelTick TickType, FActorComponentTickFunction* ThisTickFunction) override;
 	virtual FTransform GetReachableTransform() override;
+	virtual bool GetIsReadyForTraining() override;
+	virtual bool GetIsInRestState() override;
+
+	virtual void SetRobotTarget(const FTransform& TargetTransformIn) override;
+	virtual void SetRobotCurrentRewardZone(const FTransform& RewardTransformIn) override;
 
 	
-	
-	UFUNCTION(BlueprintCallable)
-	void SetTarget(const FTransform& TargetTransformIn);
 private:
 	FTransform TargetTransform;
+	bool bDropZoneIsRight = false;
 
 	//---------------------
 	//------- DEBUG -------
@@ -66,14 +69,17 @@ private:
 	FString Actuator_Jaw = FString("Jaw");
 
 	// SO100 Static Variables
-	FVector PivotOffset = FVector{-0.000030, 4.520021, 1.650041};
-	FVector JawOffset =  FVector{23, 2, 9};
+	// TODO Those values must be more precise, and probably that the way we compute the rotation is not good enough
+	// TODO Let's discuss how to improve that algorithm
+	FVector PivotOffset = FVector{-0.000030, 4.520021, 1.650041}; // From the Robot Location
+	FVector JawOffset =  FVector{23, 2, 9}; // From the Pivot
 	float MaxRange = 20.757929; // fixed_jaw_pad_3 ForwardVectorLength Delta between Rest and MaxExtend
-	float MaxYaw = 150.f;
+	float MaxYaw = 80.f;
 	
 	// Actuators Joints and Controls are expressed in doubles
-	double ClosedJaw = 0.18;
+	double ClosedJaw = -0.01;
 	double OpenedJaw = -2.0;
+	int32 JawState = 0; // 0 - Opened || 1 - Grabbing
 	
 	/**
 	 * Query the physic proxy on the RobotOwner to get the SO100 actuators values
@@ -85,9 +91,10 @@ private:
 	FSo100Actuators GetCurrentJointsFromPhysicsScene() const;
 	double GetControlJointDeltaForActuator(FString ActuatorName) const;
 	static float GetDeltaSumBetweenActuatorValues(const FSo100Actuators& A, const FSo100Actuators& B);
+	static bool AreActuatorsAlmostEqual(const FSo100Actuators& A, const FSo100Actuators& B);
 	static FSo100Actuators LerpActuators(const FSo100Actuators& A, const FSo100Actuators& B, const float Alpha); 
 	
-	// Called after every physic step
+	// Called after every physic step - this is a substep tick
 	virtual void PostPhysicStepUpdate(const float SimulationTime) override;
 	bool AnimateActuators(float SimulationTime); // Bound to the PhysicProxy post-update delegate
 	FSo100Actuators CurrentRobotActuators; // This will be updated by the post-physic delegate 
@@ -99,7 +106,8 @@ private:
 	// ----- OVER-CURRENT -----
 	// ------------------------
 	bool bDetectOverCurrent = false;
-	const float OverCurrentThreshold = 0.15;
+	const float OverCurrentThreshold = 0.1;
+	bool IsJawOverCurrent() const;
 
 	// Quick and dirty sequence of moves
 	// -1 -> Start Game, extended
@@ -111,8 +119,9 @@ private:
 	// 5 -> open jaw
 	int32 CurrentAnimationState = -1;
 	void NextAnimationState();
+	bool bBreakAfterAnimation = false;
 	
-	void BasePose();
+	void RestPose();
 	void RotateToTarget();
 	void MoveToTarget();
 	void CloseJaw();
@@ -143,7 +152,7 @@ private:
 	};
 	
 	FSo100Actuators ActuatorsDropZone {
-		PI / 2,
+		PI / 2 + 0.25,
 		-2.17,
 		0.805,
 		1.345,

Source/LuckyWorldV2/Public/Robot/RobotPawn.h

@@ -41,7 +41,6 @@ public:
 	UFUNCTION(BlueprintCallable)
 	void InitPilotComponent(); // This should have Robot type as parameter?
 
-
 	// ---------------------
 	// ------ SENSORS ------
 	// ---------------------