diff --git a/Content/Blueprint/Core/BP_3DtextHelp.uasset b/Content/Blueprint/Core/BP_3DtextHelp.uasset index 99f8f63b..3475aa92 100644 Binary files a/Content/Blueprint/Core/BP_3DtextHelp.uasset and b/Content/Blueprint/Core/BP_3DtextHelp.uasset differ diff --git a/Content/Blueprint/RobotPawnActors/BP_mujokoSO_100.uasset b/Content/Blueprint/RobotPawnActors/BP_mujokoSO_100.uasset index 30c32323..c208af64 100644 Binary files a/Content/Blueprint/RobotPawnActors/BP_mujokoSO_100.uasset and b/Content/Blueprint/RobotPawnActors/BP_mujokoSO_100.uasset differ diff --git a/Content/Levels/kitchenLevel/SM-training-so100-goal-mujoco.uasset b/Content/Levels/kitchenLevel/SM-training-so100-goal-mujoco.uasset new file mode 100644 index 00000000..36bcda6b Binary files /dev/null and b/Content/Levels/kitchenLevel/SM-training-so100-goal-mujoco.uasset differ diff --git a/Content/Levels/kitchenLevel/SM-training-so100-goal.uasset b/Content/Levels/kitchenLevel/SM-training-so100-goal.uasset new file mode 100644 index 00000000..a8f96a26 Binary files /dev/null and b/Content/Levels/kitchenLevel/SM-training-so100-goal.uasset differ diff --git a/Content/Levels/kitchenLevel/kitchenLevel1.umap b/Content/Levels/kitchenLevel/kitchenLevel1.umap index 6516b924..417621ea 100644 Binary files a/Content/Levels/kitchenLevel/kitchenLevel1.umap and b/Content/Levels/kitchenLevel/kitchenLevel1.umap differ diff --git a/Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Actors/MujocoStaticMeshActor.cpp b/Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Actors/MujocoStaticMeshActor.cpp index 823b7caf..442e9d5f 100644 --- a/Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Actors/MujocoStaticMeshActor.cpp +++ b/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(this->GetComponentByClass(UMujocoBodyComponent::StaticClass())); + MainActorBodyGeometry = Cast(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) diff --git a/Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Actors/MujocoVolumeActor.cpp b/Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Actors/MujocoVolumeActor.cpp index c3e9d9b3..46e87538 100644 --- a/Plugins/LuckyMujoco/Source/LuckyMujoco/Private/Actors/MujocoVolumeActor.cpp +++ b/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) diff --git a/Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoStaticMeshActor.h b/Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoStaticMeshActor.h index bfa0675c..d2507fb8 100644 --- a/Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoStaticMeshActor.h +++ b/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; + /** * */ @@ -17,12 +18,23 @@ public: AMujocoStaticMeshActor(); virtual void BeginPlay() override; - + + TObjectPtr MainActorBody; + FTransform GetGeometryTransform(); + +private: + TObjectPtr MainActorBodyGeometry; + + // ------------------------------ + // ------------------------------ + // -------- EDITOR ONLY --------- + // ------------------------------ + // ------------------------------ #if WITH_EDITOR +public: virtual void PreDuplicate(FObjectDuplicationParameters& DupParams) override; virtual void PostDuplicate(EDuplicateMode::Type DuplicateMode) override; #endif - TObjectPtr MainActorBody; }; diff --git a/Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoVolumeActor.h b/Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoVolumeActor.h index 831fa7b0..d9556ea4 100644 --- a/Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoVolumeActor.h +++ b/Plugins/LuckyMujoco/Source/LuckyMujoco/Public/Actors/MujocoVolumeActor.h @@ -111,8 +111,10 @@ 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 ------- // --------------------------- diff --git a/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp b/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp index 3547d80a..b3398ab4 100644 --- a/Source/LuckyWorldV2/Private/Episode/EpisodeSubSystem.cpp +++ b/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 - - // 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(DebugTextActorPtr)) + { + DebugTextActor = Cast(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()) + { + //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(CurrentRobot->RobotPilotComponent); - const auto Joints = So100PilotCmp->GetCurrentControlsFromPhysicScene(); + // const auto TimeStamp = CurrentRobot->PhysicsSceneProxy->GetMujocoData().time; + // const auto So100PilotCmp = Cast(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() + // 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()) + { + // 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); + } } diff --git a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotComponent.cpp b/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotComponent.cpp index aaa93f83..26d2e9a2 100644 --- a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotComponent.cpp +++ b/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) +{ +} diff --git a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp b/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp index aa72e25b..851d11bb 100644 --- a/Source/LuckyWorldV2/Private/Robot/PilotComponent/RobotPilotSO100Component.cpp +++ b/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() @@ -223,7 +302,8 @@ void URobotPilotSO100Component::MoveToTarget() // Get Pivot World 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("Animate -> AnimateActuators %f - %f"), AnimAlpha, GetControlJointDeltaForActuator(Actuator_Jaw)); - // UE_LOG(LogTemp, Log, TEXT("AnimationAlpha: %f - Delta: %f"), AnimAlpha, DeltaSum); - // 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 - RobotOwner->PhysicsSceneProxy->SetActuatorValue(Actuator_Jaw, FMath::Lerp( - AnimStartRobotActuators.Jaw, - AnimTargetRobotActuators.Jaw, - AnimAlpha - )); + // 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; } diff --git a/Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp b/Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp index efe16ce5..2f0848f2 100644 --- a/Source/LuckyWorldV2/Private/Robot/RobotPawn.cpp +++ b/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() diff --git a/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h b/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h index 89dae2ff..92cd95be 100644 --- a/Source/LuckyWorldV2/Public/Episode/EpisodeSubSystem.h +++ b/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,18 +18,39 @@ 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; - virtual void Tick(float DeltaTime); + // ---------------- + // ----- 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; + + // --------------------- // ------- START ------- // --------------------- /** * 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 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 CurrentRobot; - // The object that will serve for the episode - TObjectPtr 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(); + void SendEpisodeData(const FObservationPayload& Payload) const; + + + + }; diff --git a/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotComponent.h b/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotComponent.h index f4aadadd..4b610192 100644 --- a/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotComponent.h +++ b/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 diff --git a/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h b/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h index a5425621..14ae6c7c 100644 --- a/Source/LuckyWorldV2/Public/Robot/PilotComponent/RobotPilotSO100Component.h +++ b/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(); - - void BasePose(); + bool bBreakAfterAnimation = false; + + 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, diff --git a/Source/LuckyWorldV2/Public/Robot/RobotPawn.h b/Source/LuckyWorldV2/Public/Robot/RobotPawn.h index 05e7fc9e..623f08cb 100644 --- a/Source/LuckyWorldV2/Public/Robot/RobotPawn.h +++ b/Source/LuckyWorldV2/Public/Robot/RobotPawn.h @@ -40,7 +40,6 @@ public: URobotPilotComponent* RobotPilotComponent = nullptr; UFUNCTION(BlueprintCallable) void InitPilotComponent(); // This should have Robot type as parameter? - // --------------------- // ------ SENSORS ------