Index: binaries/data/mods/public/gui/session/selection_panels.js
===================================================================
--- binaries/data/mods/public/gui/session/selection_panels.js
+++ binaries/data/mods/public/gui/session/selection_panels.js
@@ -242,13 +242,13 @@
 		if (unitEntStates.some(state => !hasClass(state, "Unit")))
 			return [];
 
-		if (unitEntStates.every(state => !state.identity || !state.identity.hasSomeFormation))
+		if (unitEntStates.every(state => !state.unitAI || !state.unitAI.formations.length))
 			return [];
 
 		if (!g_AvailableFormations.has(unitEntStates[0].player))
 			g_AvailableFormations.set(unitEntStates[0].player, Engine.GuiInterfaceCall("GetAvailableFormations", unitEntStates[0].player));
 
-		return g_AvailableFormations.get(unitEntStates[0].player).filter(formation => unitEntStates.some(state => !!state.identity && state.identity.formations.indexOf(formation) != -1));
+		return g_AvailableFormations.get(unitEntStates[0].player).filter(formation => unitEntStates.some(state => !!state.unitAI && state.unitAI.formations.includes(formation)));
 	},
 	"setupButton": function(data)
 	{
Index: binaries/data/mods/public/simulation/components/GuiInterface.js
===================================================================
--- binaries/data/mods/public/simulation/components/GuiInterface.js
+++ binaries/data/mods/public/simulation/components/GuiInterface.js
@@ -276,8 +276,6 @@
 			"classes": cmpIdentity.GetClassesList(),
 			"selectionGroupName": cmpIdentity.GetSelectionGroupName(),
 			"canDelete": !cmpIdentity.IsUndeletable(),
-			"hasSomeFormation": cmpIdentity.HasSomeFormation(),
-			"formations": cmpIdentity.GetFormationsList(),
 			"controllable": cmpIdentity.IsControllable()
 		};
 
@@ -433,6 +431,7 @@
 			"canPatrol": cmpUnitAI.CanPatrol(),
 			"selectableStances": cmpUnitAI.GetSelectableStances(),
 			"isIdle": cmpUnitAI.IsIdle(),
+			"formations": cmpUnitAI.GetFormationsList(),
 			"formation": cmpUnitAI.GetFormationController()
 		};
 
Index: binaries/data/mods/public/simulation/components/Identity.js
===================================================================
--- binaries/data/mods/public/simulation/components/Identity.js
+++ binaries/data/mods/public/simulation/components/Identity.js
@@ -73,14 +73,6 @@
 		"</element>" +
 	"</optional>" +
 	"<optional>" +
-		"<element name='Formations' a:help='Optional list of space-separated formations this unit is allowed to use. Choices include: Scatter, Box, ColumnClosed, LineClosed, ColumnOpen, LineOpen, Flank, Skirmish, Wedge, Testudo, Phalanx, Syntagma, BattleLine.'>" +
-			"<attribute name='datatype'>" +
-				"<value>tokens</value>" +
-			"</attribute>" +
-			"<text/>" +
-		"</element>" +
-	"</optional>" +
-	"<optional>" +
 		"<element name='Icon'>" +
 			"<text/>" +
 		"</element>" +
@@ -111,11 +103,6 @@
 	this.controllable = this.template.Controllable ? this.template.Controllable == "true" : true;
 };
 
-Identity.prototype.HasSomeFormation = function()
-{
-	return this.GetFormationsList().length > 0;
-};
-
 Identity.prototype.GetCiv = function()
 {
 	return this.template.Civ;
@@ -164,18 +151,6 @@
 	return this.GetClassesList().indexOf(name) != -1;
 };
 
-Identity.prototype.GetFormationsList = function()
-{
-	if (this.template.Formations && this.template.Formations._string)
-		return this.template.Formations._string.split(/\s+/);
-	return [];
-};
-
-Identity.prototype.CanUseFormation = function(template)
-{
-	return this.GetFormationsList().indexOf(template) != -1;
-};
-
 Identity.prototype.GetSelectionGroupName = function()
 {
 	return this.template.SelectionGroupName || "";
Index: binaries/data/mods/public/simulation/components/UnitAI.js
===================================================================
--- binaries/data/mods/public/simulation/components/UnitAI.js
+++ binaries/data/mods/public/simulation/components/UnitAI.js
@@ -20,6 +20,14 @@
 	"<element name='FleeDistance'>" +
 		"<ref name='positiveDecimal'/>" +
 	"</element>" +
+	"<optional>" +
+		"<element name='Formations' a:help='Optional list of space-separated formations this unit is allowed to use. Choices include: Scatter, Box, ColumnClosed, LineClosed, ColumnOpen, LineOpen, Flank, Skirmish, Wedge, Testudo, Phalanx, Syntagma, BattleLine.'>" +
+			"<attribute name='datatype'>" +
+				"<value>tokens</value>" +
+			"</attribute>" +
+			"<text/>" +
+		"</element>" +
+	"</optional>" +
 	"<element name='CanGuard'>" +
 		"<data type='boolean'/>" +
 	"</element>" +
@@ -3474,6 +3482,16 @@
 	return (this.formationController != INVALID_ENTITY);
 };
 
+UnitAI.prototype.GetFormationsList = function()
+{
+	return this.template.Formations?._string?.split(/\s+/) || [];
+};
+
+UnitAI.prototype.CanUseFormation = function(formation)
+{
+	return this.GetFormationsList().includes(formation);
+};
+
 /**
  * For now, entities with a RoamDistance are animals.
  */
Index: binaries/data/mods/public/simulation/components/tests/test_GuiInterface.js
===================================================================
--- binaries/data/mods/public/simulation/components/tests/test_GuiInterface.js
+++ binaries/data/mods/public/simulation/components/tests/test_GuiInterface.js
@@ -569,8 +569,6 @@
 	"HasClass": function() { return true; },
 	"IsUndeletable": function() { return false; },
 	"IsControllable": function() { return true; },
-	"HasSomeFormation": function() { return false; },
-	"GetFormationsList": function() { return []; },
 });
 
 AddMock(10, IID_Position, {
@@ -600,8 +598,6 @@
 		"classes": ["class1", "class2"],
 		"selectionGroupName": "Selection Group Name",
 		"canDelete": true,
-		"hasSomeFormation": false,
-		"formations": [],
 		"controllable": true,
 	},
 	"position": { "x": 1, "y": 2, "z": 3 },
Index: binaries/data/mods/public/simulation/components/tests/test_Identity.js
===================================================================
--- binaries/data/mods/public/simulation/components/tests/test_Identity.js
+++ binaries/data/mods/public/simulation/components/tests/test_Identity.js
@@ -13,8 +13,6 @@
 TS_ASSERT_UNEVAL_EQUALS(cmpIdentity.GetClassesList(), []);
 TS_ASSERT_UNEVAL_EQUALS(cmpIdentity.GetVisibleClassesList(), []);
 TS_ASSERT_EQUALS(cmpIdentity.HasClass("CitizenSoldier"), false);
-TS_ASSERT_UNEVAL_EQUALS(cmpIdentity.GetFormationsList(), []);
-TS_ASSERT_EQUALS(cmpIdentity.CanUseFormation("special/formations/skirmish"), false);
 TS_ASSERT_EQUALS(cmpIdentity.GetSelectionGroupName(), "");
 TS_ASSERT_EQUALS(cmpIdentity.GetGenericName(), "Iberian Skirmisher");
 
@@ -33,7 +31,6 @@
 	"Rank": "Basic",
 	"Classes": { "_string": "CitizenSoldier Human Organic" },
 	"VisibleClasses": { "_string": "Javelineer" },
-	"Formations": { "_string": "special/formations/skirmish" },
 	"Icon": "units/iber_infantry_javelineer.png",
 	"RequiredTechnology": "phase_town"
 });
@@ -46,9 +43,6 @@
 TS_ASSERT_UNEVAL_EQUALS(cmpIdentity.GetVisibleClassesList(), ["Javelineer"]);
 TS_ASSERT_EQUALS(cmpIdentity.HasClass("CitizenSoldier"), true);
 TS_ASSERT_EQUALS(cmpIdentity.HasClass("Female"), false);
-TS_ASSERT_UNEVAL_EQUALS(cmpIdentity.GetFormationsList(), ["special/formations/skirmish"]);
-TS_ASSERT_EQUALS(cmpIdentity.CanUseFormation("special/formations/skirmish"), true);
-TS_ASSERT_EQUALS(cmpIdentity.CanUseFormation("special/formations/line"), false);
 TS_ASSERT_EQUALS(cmpIdentity.GetSelectionGroupName(), "units/iber/infantry_javelineer_b");
 TS_ASSERT_EQUALS(cmpIdentity.GetGenericName(), "Iberian Skirmisher");
 
Index: binaries/data/mods/public/simulation/helpers/Commands.js
===================================================================
--- binaries/data/mods/public/simulation/helpers/Commands.js
+++ binaries/data/mods/public/simulation/helpers/Commands.js
@@ -1597,12 +1597,11 @@
 		if (!cmpUnitAI || !cmpPosition || !cmpPosition.IsInWorld())
 			continue;
 
-		let cmpIdentity = Engine.QueryInterface(ent, IID_Identity);
 		// TODO: We only check if the formation is usable by some units
 		// if we move them to it. We should check if we can use formations
 		// for the other cases.
 		let nullFormation = (formationTemplate || cmpUnitAI.GetFormationTemplate()) == NULL_FORMATION;
-		if (nullFormation || !cmpIdentity || !cmpIdentity.CanUseFormation(formationTemplate || NULL_FORMATION))
+		if (nullFormation || !cmpUnitAI.CanUseFormation(formationTemplate || NULL_FORMATION))
 		{
 			if (nullFormation && cmpUnitAI.GetFormationController())
 				cmpUnitAI.LeaveFormation(cmd.queued || false);
@@ -1744,19 +1743,14 @@
 	// TODO: should check the player's civ is allowed to use this formation
 	// See simulation/components/Player.js GetFormations() for a list of all allowed formations
 
-	var requirements = GetFormationRequirements(formationTemplate);
+	const requirements = GetFormationRequirements(formationTemplate);
 	if (!requirements)
 		return false;
 
-	var count = 0;
-	for (let ent of ents)
-	{
-		var cmpIdentity = Engine.QueryInterface(ent, IID_Identity);
-		if (!cmpIdentity || !cmpIdentity.CanUseFormation(formationTemplate))
-			continue;
-
-		++count;
-	}
+	let count = 0;
+	for (const ent of ents)
+		if (Engine.QueryInterface(ent, IID_UnitAI)?.CanUseFormation(formationTemplate))
+			++count;
 
 	return count >= requirements.minCount;
 }