diff --git a/COMP250_1_2101327_AI/Content/Main.umap b/COMP250_1_2101327_AI/Content/Main.umap
index b2c7b41..25efbee 100644
--- a/COMP250_1_2101327_AI/Content/Main.umap
+++ b/COMP250_1_2101327_AI/Content/Main.umap
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9afc72fdb57e571cd6825ba08df57e5d348f4d0c637a1dc9e749b73856d25832
-size 40282
+oid sha256:4dabd96b91e537f957841dd9c905d313433780c6ffa20dcbad3898e6f45e6b2c
+size 40330
diff --git a/COMP250_1_2101327_AI/Content/StarterContent/Audio/Collapse01.uasset b/COMP250_1_2101327_AI/Content/StarterContent/Audio/Collapse01.uasset
index 7f0ee46..10eb1d7 100644
--- a/COMP250_1_2101327_AI/Content/StarterContent/Audio/Collapse01.uasset
+++ b/COMP250_1_2101327_AI/Content/StarterContent/Audio/Collapse01.uasset
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5565f16f5c4074ccbc5c0f305a5a21edbe258c6246cc40fd8e9f39ff7906a70b
+oid sha256:054f961ce7befe9c615be25a82ecadf95916f6b86afcb3057979e48be7bc38f5
 size 353935
diff --git a/COMP250_1_2101327_AI/Source/COMP250_1_2101327_AI/GOAP/GOAPAgent.cpp b/COMP250_1_2101327_AI/Source/COMP250_1_2101327_AI/GOAP/GOAPAgent.cpp
index 9a8e43b..0666c12 100644
--- a/COMP250_1_2101327_AI/Source/COMP250_1_2101327_AI/GOAP/GOAPAgent.cpp
+++ b/COMP250_1_2101327_AI/Source/COMP250_1_2101327_AI/GOAP/GOAPAgent.cpp
@@ -4,7 +4,6 @@
 #include "GOAPAgent.h"
 
 #include "Actions/Combo_P.h"
-#include "Actions/DefaultAttack.h"
 
 
 // Sets default values for this component's properties
@@ -23,30 +22,21 @@ void UGOAPAgent::BeginPlay()
 {
 	Super::BeginPlay();
 
-	CombatSystem = Cast<ATurnBaseCombatV2>(GetWorld()->GetGameState());
+	//Gets all classes that inherit from UGOAPAction and adds them to the AvailableActions array
+	TArray<UClass*> ActionClasses;
+	GetDerivedClasses(UGOAPAction::StaticClass(), ActionClasses, false);
+	for (UClass* ActionClass : ActionClasses)
+	{
+		UGOAPAction* Action = NewObject<UGOAPAction>(GetWorld(), ActionClass);
+		Action->Init();
+		AvailableActions.Add(Action);
+	}
 
-	UGOAPAction* DefaultAttack = NewObject<UGOAPAction>(GetWorld(), UDefaultAttack::StaticClass());
-	DefaultAttack->Init();
-	AvailableActions.Add(DefaultAttack);
-	UGOAPAction* Combo_P = NewObject<UGOAPAction>(GetWorld(), UCombo_P::StaticClass());
-	Combo_P->Init();
-	AvailableActions.Add(Combo_P);
+	//Sort AvailableActions by cost
+	AvailableActions.Sort([](const UGOAPAction& A, const UGOAPAction& B) { return A.ActionCost < B.ActionCost; });
 }
 
-
-TMap<FString, int> UGOAPAgent::GetWorldState()
-{
-	TMap<FString, int> WorldState;
-	WorldState.Add("PlayerHealth", *CombatSystem->PlayerHealth);
-	WorldState.Add("EnemyHealth", *CombatSystem->EnemyHealth);
-	WorldState.Add("ProbertiumResource", CombatSystem->EnemyProbertiumResource);
-	WorldState.Add("EisResource", CombatSystem->EnemyEisResource);
-	WorldState.Add("AzosResource", CombatSystem->EnemyAzosResource);
-	WorldState.Add("IroquoidResource", CombatSystem->EnemyIroquoidResource);
-	return WorldState;
-}
-
-bool UGOAPAgent::BuildActionGraph(UActionNode* Parent, TArray<UActionNode*>& SuccessfulLeaves, TMap<FString, int> WorldState, TMap<FString, int> AgentGoals)
+bool UGOAPAgent::BuildActionGraph(UActionNode* Parent, TArray<UActionNode*>& SuccessfulLeaves, UWorldState* WorldState, TMap<FString, int> AgentGoals)
 {
 	bool bFoundAPlan = false;
 
@@ -54,17 +44,18 @@ bool UGOAPAgent::BuildActionGraph(UActionNode* Parent, TArray<UActionNode*>& Suc
 	{
 		if (Action->CheckPreConditions(WorldState))
 		{
-			TMap<FString, int> NewWorldState = Action->ApplyEffects(WorldState);
+			UWorldState* NewWorldState = Action->ApplyEffects(WorldState);
 			UActionNode* NewNode = NewObject<UActionNode>();
 			NewNode->Init(Parent, Action, NewWorldState, Parent->Cost + Action->ActionCost);
 
-			if (NewWorldState["PlayerHealth"] <= 0)
+			if (NewWorldState->State["PlayerHealth"] <= 0)
 			{
 				SuccessfulLeaves.Add(NewNode);
 				bFoundAPlan = true;
 			}
-			else
+			else if (!SeenStates.Contains(NewWorldState) && !bFoundAPlan)
 			{
+				SeenStates.Add(NewWorldState);
 				TArray<UActionNode*> NewLeaves;
 				if (BuildActionGraph(NewNode, NewLeaves, NewWorldState, AgentGoals))
 				{
@@ -75,6 +66,17 @@ bool UGOAPAgent::BuildActionGraph(UActionNode* Parent, TArray<UActionNode*>& Suc
 		}
 	}
 
+	//Log the plan
+	if (bFoundAPlan)
+	{
+		UE_LOG(LogTemp, Warning, TEXT("Found a plan!"));
+		// UActionNode* CurrentNode = SuccessfulLeaves[0];
+		// while (CurrentNode->Parent != nullptr)
+		// {
+		// 	UE_LOG(LogTemp, Warning, TEXT("%s"), *CurrentNode->Action->GetName());
+		// 	CurrentNode = CurrentNode->Parent;
+		// }
+	}
 	return bFoundAPlan;
 }
 
@@ -86,13 +88,16 @@ void UGOAPAgent::TickComponent(float DeltaTime, ELevelTick TickType, FActorCompo
 	// ...
 }
 
-TArray<UGOAPAction*> UGOAPAgent::Plan(TMap<FString, int> WorldState, TMap<FString, int> AgentGoals)
+TArray<UGOAPAction*> UGOAPAgent::Plan(UWorldState* WorldState, TMap<FString, int> AgentGoals)
 {
 	TArray<UGOAPAction*> ActionPlan;
 	UActionNode* RootActionNode = NewObject<UActionNode>();
 	RootActionNode->Init(nullptr, nullptr, WorldState, 0);
 	TArray<UActionNode*> SuccessfulLeaves;
 
+	SeenStates.Empty();
+	SeenStates.Add(WorldState);
+
 	if (BuildActionGraph(RootActionNode, SuccessfulLeaves, WorldState, AgentGoals))
 	{
 		//Find the cheapest leaf
@@ -123,5 +128,7 @@ TArray<UGOAPAction*> UGOAPAgent::Plan(TMap<FString, int> WorldState, TMap<FStrin
 		}
 	}
 
+	//Remove the first action as it is null due to being the root node
+	ActionPlan.RemoveAt(0);
 	return ActionPlan;
 }
diff --git a/COMP250_1_2101327_AI/Source/COMP250_1_2101327_AI/GOAP/GOAPAgent.h b/COMP250_1_2101327_AI/Source/COMP250_1_2101327_AI/GOAP/GOAPAgent.h
index b1caddd..dea611f 100644
--- a/COMP250_1_2101327_AI/Source/COMP250_1_2101327_AI/GOAP/GOAPAgent.h
+++ b/COMP250_1_2101327_AI/Source/COMP250_1_2101327_AI/GOAP/GOAPAgent.h
@@ -4,8 +4,9 @@
 
 #include "CoreMinimal.h"
 #include "GOAPAction.h"
-#include "COMP250_1_2101327_AI/TurnBasedCombatV2/TurnBaseCombatV2.h"
 #include "Components/ActorComponent.h"
+#include "Containers/UnrealString.h"
+#include "WorldState.h"
 #include "GOAPAgent.generated.h"
 
 UCLASS()
@@ -16,10 +17,10 @@ class UActionNode : public UObject
 public:
 	UActionNode* Parent;
 	UGOAPAction* Action;
-	TMap<FString, int> State;
+	UWorldState* State;
 	float Cost;
 
-	void Init(UActionNode* ParentNode, UGOAPAction* GOAPAction, TMap<FString, int> ProceduralState, float TotalCost)
+	void Init(UActionNode* ParentNode, UGOAPAction* GOAPAction, UWorldState* ProceduralState, float TotalCost)
 	{
 		Parent = ParentNode;
 		Action = GOAPAction;
@@ -50,19 +51,14 @@ protected:
 	// Called when the game starts
 	virtual void BeginPlay() override;
 
-	UPROPERTY()
-	ATurnBaseCombatV2* CombatSystem;
-
 	UFUNCTION()
-	bool BuildActionGraph(UActionNode* Parent, TArray<UActionNode*>& SuccessfulLeaves, TMap<FString, int> WorldState, TMap<FString, int> AgentGoals);
+	bool BuildActionGraph(UActionNode* Parent, TArray<UActionNode*>& SuccessfulLeaves, UWorldState* WorldState, TMap<FString, int> AgentGoals);
 
+	TSet<UWorldState*> SeenStates;
 public:
 	// Called every frame
 	virtual void TickComponent(float DeltaTime, ELevelTick TickType, FActorComponentTickFunction* ThisTickFunction) override;
 
 	UFUNCTION()
-	TArray<UGOAPAction*> Plan(TMap<FString, int> WorldState, TMap<FString, int> AgentGoals);
-
-	UFUNCTION()
-	TMap<FString, int> GetWorldState();
+	TArray<UGOAPAction*> Plan(UWorldState* WorldState, TMap<FString, int> AgentGoals);
 };