Merge pull request 'Patch MuJoCo plugin' (#50) from ft-pre-post-substep-delegates into main

Reviewed-on: #50
Reviewed-by: evanvyktori <evan@luckyrobots.com>

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

@@ -109,9 +109,21 @@ template <typename ComponentType> void AMujocoVolumeActor::AssignComponentsToArr
 	}
 }
 
+template <typename UserClass>
+void AMujocoVolumeActor::BindPostPhysicDelegate(UserClass* Object, void(UserClass::* Func)(float))
+{
+	PostPhysicUpdateDelegate.BindUObject(Object, Func);
+}
+
+
 void AMujocoVolumeActor::InitializeMujoco()
 {
-	if (!Options)
+	InitializeMujocoScene_WithContactExclusion(TMap<FString, FString>{});
+}
+
+void AMujocoVolumeActor::InitializeMujocoScene_WithContactExclusion(const TMap<FString, FString>& ContactExclusion)
+{
+		if (!Options)
 	{
 		return;
 	}
@@ -149,7 +161,7 @@ void AMujocoVolumeActor::InitializeMujoco()
 		}
 
 		TUniquePtr<FMujocoXmlGenerator> Generator = MakeUnique<FMujocoXmlGenerator>();
-		TUniquePtr<tinyxml2::XMLDocument> Doc = Generator->GenerateMujocoXml(Options, Objects, ExportFilename);
+		TUniquePtr<tinyxml2::XMLDocument> Doc = Generator->GenerateMujocoXml(Options, Objects, ExportFilename, ContactExclusion);
 
 		FString XmlString;
 		tinyxml2::XMLPrinter Printer;
@@ -164,6 +176,8 @@ void AMujocoVolumeActor::InitializeMujoco()
 			return;
 		}
 
+		// TODO Here we should check for mujoco dll if we are on windows and LOG an error if there is not + gently quit the game
+		
 		std::array<char, 1024> ErrMsg{};
 		MujocoModel = MakeMujocoModelPtr(mj_loadXML(TCHAR_TO_ANSI(*ExportFilename), nullptr, ErrMsg.data(), ErrMsg.size()));
 		if (!MujocoModel)
@@ -196,6 +210,12 @@ void AMujocoVolumeActor::InitializeMujoco()
 	}
 }
 
+mjData_& AMujocoVolumeActor::GetMujocoData() const
+{
+	check(MujocoData.IsValid());
+	return *MujocoData.Get();
+}
+
 void AMujocoVolumeActor::SetActuatorValue(const FString& ActuatorName, double Value)
 {
 	if (MujocoModel)
@@ -391,11 +411,13 @@ void AMujocoVolumeActor::Tick(float DeltaTime)
 {
 	if (MujocoData)
 	{
-		mj_step(MujocoModel.Get(), MujocoData.Get());
+		// Default SimStepsPerGameFrame = 16 x 1ms for roughly 60FPS
-
+		// If ticks are strictly 16.666 ms then simulation will lag behind UE clock
-		for (int32 Frame = 0; Frame < FrameSkip; ++Frame)
+		// Anyway, we need a proper synchronisation between UE and Sim times - probably using TempoTime?
+		for (int32 SimStep = 0; SimStep < SimStepsPerGameFrame; ++SimStep)
 		{
 			mj_step(MujocoModel.Get(), MujocoData.Get());
+			PostPhysicUpdateDelegate.ExecuteIfBound(MujocoData->time);
 		}
 
 		for (int32 i = 1; i < BodyComponents.Num(); ++i)

Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Misc/MujocoXmlGenerator.cpp

@@ -1255,7 +1255,11 @@ bool FMujocoXmlGenerator::ParseActorAsset(AActor* Actor, tinyxml2::XMLElement* R
 	return true;
 }
 
-TUniquePtr<tinyxml2::XMLDocument> FMujocoXmlGenerator::GenerateMujocoXml(const TObjectPtr<UMujocoExportOptions>& ExportOptions, const TArray<UObject*>& Objects, const FString& ExportFilename)
+TUniquePtr<tinyxml2::XMLDocument> FMujocoXmlGenerator::GenerateMujocoXml(
+	const TObjectPtr<UMujocoExportOptions>& ExportOptions,
+	const TArray<UObject*>& Objects,
+	const FString& ExportFilename,
+	TMap<FString, FString> ContactExclusion)
 {
 	TUniquePtr<tinyxml2::XMLDocument> Doc = MakeUnique<tinyxml2::XMLDocument>();
 	Doc->InsertFirstChild(Doc->NewDeclaration("xml version=\"1.0\" encoding=\"utf-8\""));
@@ -1268,6 +1272,30 @@ TUniquePtr<tinyxml2::XMLDocument> FMujocoXmlGenerator::GenerateMujocoXml(const T
 	Root->InsertEndChild(Doc->NewElement("equality"));
 	Root->InsertEndChild(Doc->NewElement("tendon"));
 	Root->InsertEndChild(Doc->NewElement("actuator"));
+
+	// TODO Refacto using property from the MujocoActor aka the robot
+	// TODO Maybe using a ParentObject as parameter which holds that information?
+	// TODO Need actor refactoring anyway - Merge Pawn and Mujoco Actor?
+	if (!ContactExclusion.IsEmpty())
+	{
+		Root->InsertEndChild(Doc->NewElement("contact"));
+
+		for (auto& [Body1, Body2] : ContactExclusion)
+		{
+			// Don't add empty strings to the map please!
+			if (Body1.IsEmpty() || Body2.IsEmpty()) continue;
+
+			// Create the contact exclusion - formatted as (SO100 example)
+			// <exclude body1="Body_Base" body2="Body_Rotation_Pitch"/> 
+			const auto Exclude = Doc->NewElement("exclude");
+			Exclude->SetAttribute("body1", TCHAR_TO_ANSI(*Body1));
+			Exclude->SetAttribute("body2", TCHAR_TO_ANSI(*Body2));
+
+			// Add exclusion to the contact tag
+			Root->FirstChildElement("contact")->InsertEndChild(Exclude); 
+		}
+	}
+	
 	if (ExportOptions->bAddSkyBox)
 	{
 		tinyxml2::XMLElement* TextureElement = AssetRoot->GetDocument()->NewElement("texture");

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

@@ -21,11 +21,44 @@ DECLARE_DYNAMIC_MULTICAST_DELEGATE(FOnMujocoCompileBegin);
 DECLARE_DYNAMIC_MULTICAST_DELEGATE_OneParam(FOnMujocoCompileError, FString, Error);
 DECLARE_DYNAMIC_MULTICAST_DELEGATE(FOnMujocoCompileSuccess);
 
+// To be called after mj_step
+DECLARE_DELEGATE_OneParam(FPostPhysicUpdate, float);
+
+
 UCLASS(Blueprintable, BlueprintType)
 class LUCKYMUJOCO_API AMujocoVolumeActor : public AActor
 {
 	GENERATED_BODY()
 
+public:
+	AMujocoVolumeActor();
+
+
+	/**
+	 * Initialize the sim scene in headless mujoco
+	 */
+	UFUNCTION(BlueprintCallable, Category = "Mujoco")
+	void InitializeMujoco();
+
+
+	/**
+	* Initialize the sim scene in headless mujoco with a list of contact exclusion
+	* @param ContactExclusion a list of pairs that should be patched in the xml file for contact exclusion (no friction, no collision)
+	* TODO Can't use a default empty map as parameter in blueprints? We shouldn't need to have 2 functions 
+	* TODO ContactExclusion should be stored as a property of MujocoActor and not passed in the scene init
+	* TODO This require to cast the Actor in addition to the components list
+	*/
+	UFUNCTION(BlueprintCallable, Category = "Mujoco")
+	void InitializeMujocoScene_WithContactExclusion(const TMap<FString, FString>& ContactExclusion);
+
+protected:
+	virtual void BeginPlay() override;
+	virtual void EndPlay(const EEndPlayReason::Type EndPlayReason) override;
+	virtual void PostRegisterAllComponents() override;	
+	virtual void Tick(float DeltaTime) override;
+	virtual void PostInitializeComponents() override;
+
+private:
 	TMujocoModelPtr MujocoModel;
 	TMujocoDataPtr MujocoData;
 
@@ -50,16 +83,11 @@ class LUCKYMUJOCO_API AMujocoVolumeActor : public AActor
 	UPROPERTY(Transient, VisibleAnywhere, Category = "Mujoco | Debug")
 	TArray<TSoftObjectPtr<UMujocoTendonComponent>> TendonComponents;	
 
-	UFUNCTION(BlueprintCallable, Category = "Mujoco")
-	void InitializeMujoco();
-
 	template <typename ComponentType> void AssignComponentsToArray(UWorld* World, TArray<TSoftObjectPtr<ComponentType>>& ComponentArray);
 
 public:
-	AMujocoVolumeActor();	
+	UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = "Mujoco | Simulation", meta = (Min = 1, Max = 100, ClampMin = 0, ClampMax = 100))
-
+	int32 SimStepsPerGameFrame = 16;
-	UPROPERTY(EditAnywhere, BlueprintReadWrite, Category = "Mujoco | Simulation", meta = (Min = 0, Max = 100, ClampMin = 0, ClampMax = 100))
-	int32 FrameSkip = 0;
 
 	UPROPERTY(EditDefaultsOnly, BlueprintReadWrite, Category = "Mujoco")
 	TObjectPtr<UBillboardComponent> SpriteComponent;
@@ -76,6 +104,30 @@ public:
 	UPROPERTY(BlueprintAssignable, Category = "Mujoco | Events")
 	FOnMujocoCompileSuccess OnMujocoCompileSuccess;
 
+	/**
+	 * Access MuJoCo scene options and data (e.g. Simulation time) - This is mutable, be careful
+	 * @return mjData_ - Full access to mujoco scene options and data 
+	 */
+	mjData_& GetMujocoData() const;
+	
+	// ---------------------------
+	// ------- POST UPDATE -------
+	// ---------------------------
+private:
+	FPostPhysicUpdate PostPhysicUpdateDelegate;
+public:	
+	/**
+	 * Register a delegate to be executed after mj_step, useful to fine control actuators 
+	 * @param Object - The Class to call
+	 * @param Func - The Function to call - takes a float as parameter which is the current simulation timestamp
+	 */
+	template<typename UserClass>
+	void BindPostPhysicDelegate(UserClass* Object, void (UserClass::*Func)(float));
+	
+
+	// -------------------------
+	// ------- ACTUATORS -------
+	// -------------------------
 	UFUNCTION(BlueprintCallable, Category = "Mujoco")
 	void SetActuatorValue(const FString& ActuatorName, double Value);
 
@@ -94,6 +146,9 @@ public:
 	UFUNCTION(BlueprintCallable, BlueprintPure, Category = "Mujoco")
 	FVector2D GetActuatorRangeByIndex(int32 ActuatorIndex) const;
 
+	// ----------------------
+	// ------- JOINTS -------
+	// ----------------------
 	UFUNCTION(BlueprintCallable, Category = "Mujoco")
 	void SetJointValue(const FString& JointName, double Value);
 
@@ -105,16 +160,4 @@ public:
 
 	UFUNCTION(BlueprintCallable, BlueprintPure, Category = "Mujoco")
 	double GetJointValueByIndex(int32 JointIndex) const;
-
-	
-	virtual void PostRegisterAllComponents() override;	
-
-	virtual void Tick(float DeltaTime) override;
-
-	virtual void PostInitializeComponents() override;
-
-protected:
-
-	virtual void BeginPlay() override;
-	virtual void EndPlay(const EEndPlayReason::Type EndPlayReason) override;
 };

Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Misc/MujocoOptions.h

@@ -307,13 +307,13 @@ struct LUCKYMUJOCO_API FMujocoOptions
 {
 	GENERATED_BODY()
 
-	/** Override TimeStep setting. When enabled, the default value (0.002) is applied. */
+	/** Override TimeStep setting. When enabled, the default value (0.001) is applied. */
 	UPROPERTY(config, EditAnywhere, BlueprintReadWrite, Category = "Mujoco", meta = (InlineEditConditionToggle))
 	bool bOverrideTimeStep = false;
 
 	/** Simulation time step in seconds. */
 	UPROPERTY(config, EditAnywhere, BlueprintReadWrite, Category = "Mujoco", meta = (Attribute = "timestep", EditCondition = "bOverrideTimeStep", ClampMin = 0.0001, ClampMax = 0.01, UIMin = 0.0001, UIMax = 0.01))
-	float TimeStep = 0.002f;
+	float TimeStep = 0.001f; // Default to 1ms
 
 	/** Override API Rate setting. */
 	UPROPERTY(config, EditAnywhere, BlueprintReadWrite, Category = "Mujoco", meta = (InlineEditConditionToggle))

Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Misc/MujocoXmlGenerator.h

@@ -36,5 +36,18 @@ class LUCKYMUJOCO_API FMujocoXmlGenerator
 	TMap<FString, TSoftObjectPtr<UObject>> ObjectMap;
 
 public:
-	TUniquePtr<tinyxml2::XMLDocument> GenerateMujocoXml(const TObjectPtr<UMujocoExportOptions>& ExportOptions, const TArray<UObject*>& Objects, const FString& ExportFilename);
+
+	/**
+	 * Generate the XML file used by the MuJoCo headless scene representing the physics simulation
+	 * @param ExportOptions TODO - What they can be?
+	 * @param Objects Actors and Components to be added to the scene
+	 * @param ExportFilename pretty much that 
+	 * @param ContactExclusion a list of pairs that should be patched in the xml file for contact exclusion (no friction, no collision)
+	 * TODO ContactExclusion should be stored as a property of MujocoActor and not passed in the scene init
+	 */
+	TUniquePtr<tinyxml2::XMLDocument> GenerateMujocoXml(
+		const TObjectPtr<UMujocoExportOptions>& ExportOptions,
+		const TArray<UObject*>& Objects,
+		const FString& ExportFilename,
+		TMap<FString, FString> ContactExclusion);
 };

Plugins/LuckyMujoco/Source/LuckyMujocoEditor/Private/MujocoExporter.cpp

@@ -141,8 +141,8 @@ public:
 			CancelButton->SetEnabled(false);
 			RequestDestroyWindow();
 
-			TUniquePtr<FMujocoXmlGenerator> Generator = MakeUnique<FMujocoXmlGenerator>();			
+			const TUniquePtr<FMujocoXmlGenerator> Generator = MakeUnique<FMujocoXmlGenerator>();
-			TUniquePtr<tinyxml2::XMLDocument> Doc = Generator->GenerateMujocoXml(ExportOptions, Objects, ExportFilename);
+			const TUniquePtr<tinyxml2::XMLDocument> Doc = Generator->GenerateMujocoXml(ExportOptions, Objects, ExportFilename, TMap<FString, FString>{});
 
 			FString XmlString;
 			tinyxml2::XMLPrinter Printer;