UNSTBL - Attempt to fix So100 training

+ Teleport (success) the robot arm to the training zone
+ Add geometries to the physics scene (success)
+ Target Training Object by name (success)
- Teleport geometry to target zone is broken

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

@@ -9,6 +9,7 @@
 #include "Components/MujocoTendonComponent.h"
 #include "array"
 #include "LuckyMujoco.h"
+#include "Actors/MujocoStaticMeshActor.h"
 
 AMujocoVolumeActor::AMujocoVolumeActor()
 {
@@ -146,20 +147,18 @@ void AMujocoVolumeActor::InitializeMujocoScene_Internal()
 		TArray<UObject*> Objects;
 
 		// Old method, taking the components from the scene - this is bad because our description of the robot is imperfect
-		for (AActor* Actor : World->GetCurrentLevel()->Actors)
+		const auto Actors = World->GetCurrentLevel()->Actors; 
+		for (auto& Actor : Actors)
 		{
-			if (!Actor || Actor->IsPendingKillPending())
-			{
-				continue;
-			}
+			// Actor is invalid
+			if (!Actor || Actor->IsPendingKillPending()) continue;
 
 			TInlineComponentArray<UMujocoBodyComponent*> MujocoBodyComponents;
 			Actor->GetComponents<UMujocoBodyComponent>(MujocoBodyComponents);
-			if (MujocoBodyComponents.Num() == 0)
-			{
-				continue;
-			}
-			if (!bUseRobotModelInclude)
+			
+			// Even if we don't include robot model we still want objects of the UE scene in the physics scene
+			const auto MujocoStaticMeshActor = Cast<AMujocoStaticMeshActor>(Actor); 
+			if ((!MujocoBodyComponents.IsEmpty() && !bUseRobotModelInclude) || MujocoStaticMeshActor)
 			{
 				Objects.Add(Actor);
 			}
@@ -311,6 +310,41 @@ FTransform AMujocoVolumeActor::GetGeometryTransform(const FString& BodyName) con
 	return FTransform(Rotation, Position);
 }
 
+void AMujocoVolumeActor::TeleportRootBody(const FString& BodyName, const FTransform& NewTransform)
+{
+	if (!MujocoModel.IsValid() || !MujocoData.IsValid()) return;
+
+	const int32 BodyID = mj_name2id(MujocoModel.Get(), mjOBJ_BODY, TCHAR_TO_ANSI(*BodyName));
+	if (BodyID < 0) {
+		UE_LOG(LogTemp, Error, TEXT("Teleport failed: Body not found: %s"), *BodyName);
+		return;
+	}
+
+	// Step 1: Set new position (convert Unreal → MuJoCo: cm to m and Y-flip)
+	const FVector NewLocation = NewTransform.GetLocation() / 100.f; // cm → m
+	MujocoModel->body_pos[3 * BodyID + 0] = NewLocation.X;
+	MujocoModel->body_pos[3 * BodyID + 1] = -NewLocation.Y; // flip Y for Unreal → MJ
+	MujocoModel->body_pos[3 * BodyID + 2] = NewLocation.Z;
+
+	// Step 2: Set new orientation (Unreal XYZW → MuJoCo WXYZ, and flip Y/Z)
+	const FQuat UnrealQuat = NewTransform.GetRotation();
+	mjtNum NewQuat[4] = {
+		UnrealQuat.W,
+		UnrealQuat.X,
+		-UnrealQuat.Y,
+		-UnrealQuat.Z
+	};
+	mju_normalize4(NewQuat);
+	mju_copy(&MujocoModel->body_quat[4 * BodyID], NewQuat, 4);
+
+	// Step 3: Reset and re-run physics pipeline
+	mj_resetData(MujocoModel.Get(), MujocoData.Get());
+	mj_forward(MujocoModel.Get(), MujocoData.Get());
+
+	UE_LOG(LogTemp, Log, TEXT("Teleported body '%s' to (%f, %f, %f) [UE_coords=cm]"), *BodyName, NewLocation.X, NewLocation.Y, NewLocation.Z);
+}
+
+
 void AMujocoVolumeActor::SetActuatorValue(const FString& ActuatorName, double Value)
 {
 	if (MujocoModel)

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

@@ -124,11 +124,19 @@ public:
 	 */
 	mjData_& GetMujocoData() const;
 
+	/**
+	 * Updates a free joint geometry in the physics scene
+	 * @param BodyName 
+	 * @param NewTransform 
+	 */
 	UFUNCTION(BlueprintCallable, Category = "Mujoco")
 	void UpdateGeomTransform(const FString& BodyName, const FTransform& NewTransform);
 
 	FTransform GetGeometryTransform(const FString& BodyName) const;
-	
+
+	UFUNCTION(BlueprintCallable, Category = "Mujoco")
+	void TeleportRootBody(const FString& BodyName, const FTransform& NewTransform);
+
 	// ---------------------------
 	// ------- POST UPDATE -------
 	// ---------------------------

Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp

@@ -97,7 +97,11 @@ void UEpisodeSubSystem::UpdateDebugTextActor() const
 	TextRender->SetText(FText::FromString(Txt));
 }
 
-void UEpisodeSubSystem::StartTraining(const int32 EpisodesCountIn, FString BaseImageDataPathIn, FString TaskDescriptionIn)
+void UEpisodeSubSystem::StartTraining(
+	const int32 EpisodesCountIn,
+	const FString& BaseImageDataPathIn,
+	const FString& TaskDescriptionIn,
+	const FString& TrainingObjectName)
 {
 	// Debug
 	const auto DebugTextActorPtr = UGameplayStatics::GetActorOfClass(this->GetWorld(), ATextRenderActor::StaticClass());
@@ -107,7 +111,7 @@ void UEpisodeSubSystem::StartTraining(const int32 EpisodesCountIn, FString BaseI
 	}
 	
 	// Robot and Exercise
-	FindEpisodeObjectFromScene();
+	FindTrainingObjectFromScene(TrainingObjectName);
 	FindRobotPawnFromScene();
 	EpisodesToCapture = EpisodesCountIn;
 	SuccessEpisodes = 0;
@@ -207,13 +211,15 @@ bool UEpisodeSubSystem::CheckEpisodeCompletion()
 	return true;
 }
 
-void UEpisodeSubSystem::FindEpisodeObjectFromScene()
+void UEpisodeSubSystem::FindTrainingObjectFromScene(const FString& TrainingObjectName)
 {
 	TArray<AActor*> MujocoObjects;
 	UGameplayStatics::GetAllActorsOfClass(this->GetWorld(), AMujocoStaticMeshActor::StaticClass(), MujocoObjects);
-	if (MujocoObjects.IsValidIndex(0) && Cast<AMujocoStaticMeshActor>(MujocoObjects[0]))
+	for (const auto MujocoActor : MujocoObjects)
 	{
-		EpisodeTargetObject = Cast<AMujocoStaticMeshActor>(MujocoObjects[0]);		
+		const auto ActorName = MujocoActor->GetActorNameOrLabel();
+		if (ActorName != TrainingObjectName) continue;
+		EpisodeTargetObject = Cast<AMujocoStaticMeshActor>(MujocoActor);		
 	}
 }
 

Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp

@@ -5,6 +5,7 @@
 #include "Misc/Paths.h"
 #include "Robot/PilotComponent/RobotPilotMultiRotorDrone.h"
 #include "Robot/So100/RobotPilotCmpSo100.h"
+#include "_Utils/FileUtils.h"
 
 ARobotPawn::ARobotPawn()
 {
@@ -21,20 +22,25 @@ void ARobotPawn::InitRobot()
 	InitPilotComponent();
 }
 
+
+
+
 FString ARobotPawn::GetMenagerieXmlFullFilePath() const
 {
-	FString RobotMenageriePath = FPaths::ProjectDir() / "mujoco_menagerie-main" / RobotRelativeBasePath;
-	FString RobotMenagerieXmlPath = RobotMenageriePath / MenagerieXmlFile;  
-	
+	const FString RobotMenageriePath = FPaths::ProjectDir() / "mujoco_menagerie-main" / RobotRelativeBasePath;
+	const FString RobotMenagerieXmlPath = RobotMenageriePath / MenagerieXmlFile;  
+
 	IPlatformFile& PlatformFile = FPlatformFileManager::Get().GetPlatformFile();
     
 	// Ensure destination directory exists
-	FString DestDir = FPaths::GetPath(RobotMenageriePath / "assets");
+	const FString DestDir = FPaths::GetPath(RobotMenageriePath / "assets");
 	if (!PlatformFile.DirectoryExists(*DestDir)) return RobotMenageriePath / MenagerieXmlFile;
-	
-	// Copy the file
-	FString DestPath = RobotMenageriePath / "assets" / MenagerieXmlFile; 
-	PlatformFile.CopyFile(*DestPath, *RobotMenagerieXmlPath);
+
+	FString DestPath = RobotMenageriePath / "assets" / MenagerieXmlFile;
+
+	// TODO This is an ugly fix, we should be able to load mesh assets in a better way
+	// TODO and we should not need to strip keyframes, or at least understand why we need to do so 
+	UFileUtils::CopyAndStripKeyframes(RobotMenagerieXmlPath, DestPath);
 	
 	return  DestPath;
 }

Source/LuckyWorldV2/Private/_Utils/FileUtils.cpp

@@ -6,8 +6,6 @@ UFileUtils::UFileUtils()
 {
 }
 
-
-
 bool UFileUtils::WriteJsonlFile(const TArray<FString>& JsonLines, const FString& BasePath, const FString& FileName)
 {
 	// Ensure the directory exists
@@ -21,4 +19,44 @@ bool UFileUtils::WriteJsonlFile(const TArray<FString>& JsonLines, const FString&
 
 	// Write to file
 	return FFileHelper::SaveStringToFile(FileContent, *FullFilePath);
+}
+
+bool UFileUtils::CopyAndStripKeyframes(const FString& SourcePath, const FString& DestPath)
+{
+	FString FileContents;
+
+	// Load source XML file
+	if (!FFileHelper::LoadFileToString(FileContents, *SourcePath))
+	{
+		UE_LOG(LogTemp, Error, TEXT("Failed to read XML file: %s"), *SourcePath);
+		return false;
+	}
+
+	// Regex to remove all <keyframe>...</keyframe> blocks
+	FRegexPattern Pattern(TEXT("<keyframe\\b[^>]*>[\\s\\S]*?<\\/keyframe>"));
+	FRegexMatcher Matcher(Pattern, FileContents);
+
+	while (Matcher.FindNext())
+	{
+		FileContents.RemoveAt(Matcher.GetMatchBeginning(), Matcher.GetMatchEnding() - Matcher.GetMatchBeginning());
+		Matcher = FRegexMatcher(Pattern, FileContents); // reset matcher after modification
+	}
+
+	// Ensure destination folder exists
+	const FString DestDir = FPaths::GetPath(DestPath);
+	IPlatformFile& PlatformFile = FPlatformFileManager::Get().GetPlatformFile();
+	if (!PlatformFile.DirectoryExists(*DestDir))
+	{
+		PlatformFile.CreateDirectoryTree(*DestDir);
+	}
+
+	// Save cleaned XML
+	if (!FFileHelper::SaveStringToFile(FileContents, *DestPath))
+	{
+		UE_LOG(LogTemp, Error, TEXT("Failed to write cleaned XML file: %s"), *DestPath);
+		return false;
+	}
+
+	UE_LOG(LogTemp, Log, TEXT("Copied and stripped keyframes: %s -> %s"), *SourcePath, *DestPath);
+	return true;
 }

Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h

@@ -70,7 +70,8 @@ public:
 	 * Called by the UI when pressing the "Capture" button
 	 */
 	UFUNCTION(BlueprintCallable)
-	void StartTraining(int32 EpisodesCountIn, FString BaseImageDataPathIn, FString TaskDescriptionIn);
+	void StartTraining(int32 EpisodesCountIn, const FString& BaseImageDataPathIn, const FString& TaskDescriptionIn, const FString&
+	                   TrainingObjectName);
 	void EndTraining();
 
 	UPROPERTY()
@@ -103,7 +104,7 @@ private:
 	int32 EpisodesToCapture = 0;
 	int32 CapturedEpisodes = 0;
 	
-	void FindEpisodeObjectFromScene();
+	void FindTrainingObjectFromScene(const FString& TrainingObjectName);
 	void FindRobotPawnFromScene();
 	
 	UPROPERTY()

Source/LuckyWorldV2/Public/Robot/RobotPawn.h

@@ -22,7 +22,8 @@ public:
 	// TODO Called by GameInstance after robot has been spawned
 	UFUNCTION(BlueprintCallable)
 	void InitRobot();
-	
+	static bool CopyAndStripKeyframes(const FString& SourcePath, const FString& DestPath);
+
 	UPROPERTY(EditAnywhere, BlueprintReadWrite)
 	ERobotsName RobotType = ERobotsName::None; // This value must be set in the pawn
 

Source/LuckyWorldV2/Public/Robot/So100/RobotPilotCmpSo100.h

@@ -99,8 +99,8 @@ private:
 	float MaxYaw = 80.f;
 	
 	// Actuators Joints and Controls are expressed in doubles
-	double ClosedJaw = -0.01;
-	double OpenedJaw = -2.0;
+	double ClosedJaw = 0.01;
+	double OpenedJaw = 1.5;
 	int32 JawState = 0; // 0 - Opened || 1 - Grabbing
 	
 	/**
@@ -164,7 +164,7 @@ private:
 		3.105,
 		-1.5,
 		1.47,
-		-1.39
+		OpenedJaw
 	};
 
 	FSo100Actuators ActuatorsMaxExtendPosition {
@@ -173,7 +173,7 @@ private:
 		-0.01707,
 		-0.075,
 		1.469020,
-		-1.389073
+		0
 	};
 	
 	FSo100Actuators ActuatorsDropZone {

Source/LuckyWorldV2/Public/_Utils/FileUtils.h

@@ -11,4 +11,5 @@ public:
 	UFileUtils();
 
 	static bool WriteJsonlFile(const TArray<FString>& JsonLines, const FString& BasePath, const FString& FileName);
+	static bool CopyAndStripKeyframes(const FString& SourcePath, const FString& DestPath);
 };