cogames-agents 0.0.0.7__cp312-cp312-macosx_11_0_arm64.whl

cogames_agents/policy/nim_agents/cogsguard_align_all_agents.nim

@@ -0,0 +1,569 @@
+import
+  std/[random, sets, tables, options, strutils],
+  fidget2/measure,
+  common
+
+const
+  DepotTags = ["junction", "junction", "supply_depot"]
+  HubTags = ["hub", "hub", "main_nexus"]
+  StationTags = ["aligner_station", "scrambler_station", "miner_station", "scout_station"]
+  ResourceNames = ["carbon", "oxygen", "germanium", "silicon"]
+  ExploreSteps = 8
+  DebugEnabled = false
+  DebugEvery = 50
+  DebugScramblerId = 0
+  DebugAlignerId = 1
+  DebugMinerId = 2
+
+const Offsets4 = [
+  Location(x: 1, y: 0),
+  Location(x: 0, y: 1),
+  Location(x: -1, y: 0),
+  Location(x: 0, y: -1),
+]
+
+type
+  PendingAction = enum
+    PendingNone = 0,
+    PendingAlign = 1,
+    PendingScramble = 2
+
+  CogsguardAlignAllAgent* = ref object
+    agentId*: int
+    cfg: Config
+    random: Rand
+    map: Table[Location, seq[FeatureValue]]
+    seen: HashSet[Location]
+    unreachables: HashSet[Location]
+    location: Location
+    bump: bool
+
+    exploreDirIndex: int
+    exploreSteps: int
+
+    stations: Table[string, Location]
+    depots: Table[Location, int] # -1 clips, 0 neutral/unknown, 1 cogs
+    extractors: Table[string, seq[Location]]
+    extractorRemaining: Table[Location, int]
+    hub: Option[Location]
+    chest: Option[Location]
+    actionIds: Table[string, int]
+
+    pendingAction: PendingAction
+    pendingTarget: Option[Location]
+    lastHeart: int
+    stepCount: int
+
+  CogsguardAlignAllPolicy* = ref object
+    agents*: seq[CogsguardAlignAllAgent]
+
+proc getVibeName(agent: CogsguardAlignAllAgent, vibeId: int): string =
+  if vibeId >= 0 and vibeId < agent.cfg.vibeNames.len:
+    return agent.cfg.vibeNames[vibeId]
+  return "default"
+
+proc getActionId(agent: CogsguardAlignAllAgent, name: string): int =
+  return agent.actionIds.getOrDefault(name, agent.cfg.actions.noop)
+
+proc actionForVibe(agent: CogsguardAlignAllAgent, vibe: string): int =
+  let actionName = "change_vibe_" & vibe
+  return agent.getActionId(actionName)
+
+proc featureValue(
+  features: seq[FeatureValue],
+  featureId: int
+): int =
+  for feature in features:
+    if feature.featureId == featureId:
+      return feature.value
+  return -1
+
+proc getTagNames(cfg: Config, features: seq[FeatureValue]): HashSet[string] =
+  result = initHashSet[string]()
+  for feature in features:
+    if feature.featureId == cfg.features.tag:
+      if feature.value >= 0 and feature.value < cfg.config.tags.len:
+        result.incl(cfg.config.tags[feature.value])
+
+proc getAlignment(tagNames: HashSet[string], clipped: int): int =
+  if "collective:cogs" in tagNames or "cogs" in tagNames:
+    return 1
+  if "collective:clips" in tagNames or "clips" in tagNames:
+    return -1
+  if clipped > 0:
+    return -1
+  return 0
+
+proc isResourceExtractor(tagName: string): bool =
+  for resource in ResourceNames:
+    if resource & "_extractor" in tagName:
+      return true
+    if resource & "_chest" in tagName:
+      return true
+  return false
+
+proc updateDiscoveries(agent: CogsguardAlignAllAgent, visible: Table[Location, seq[FeatureValue]]) =
+  for location, features in visible:
+    let tagNames = getTagNames(agent.cfg, features)
+    if tagNames.len == 0:
+      continue
+
+    let absoluteLoc = Location(x: location.x + agent.location.x, y: location.y + agent.location.y)
+    let clipped = featureValue(features, agent.cfg.features.clipped)
+    let alignment = getAlignment(tagNames, clipped)
+
+    for tagName in tagNames.items:
+      for stationName in StationTags.items:
+        if tagName == stationName:
+          agent.stations[stationName] = absoluteLoc
+      if tagName in HubTags:
+        agent.hub = some(absoluteLoc)
+      if tagName in DepotTags:
+        agent.depots[absoluteLoc] = alignment
+
+      if tagName == "chest" and not tagName.isResourceExtractor():
+        agent.chest = some(absoluteLoc)
+
+      if tagName.isResourceExtractor():
+        for resource in ResourceNames:
+          if resource & "_extractor" in tagName or resource & "_chest" in tagName:
+            var locations = agent.extractors.getOrDefault(resource, @[])
+            if absoluteLoc notin locations:
+              locations.add(absoluteLoc)
+            agent.extractors[resource] = locations
+
+            let remaining = featureValue(features, agent.cfg.features.remainingUses)
+            if remaining != -1:
+              agent.extractorRemaining[absoluteLoc] = remaining
+
+proc updateMap(agent: CogsguardAlignAllAgent, visible: Table[Location, seq[FeatureValue]]) {.measure.} =
+  if agent.map.len == 0:
+    agent.map = visible
+    agent.location = Location(x: 0, y: 0)
+    for x in -5 .. 5:
+      for y in -5 .. 5:
+        agent.seen.incl(Location(x: x, y: y))
+    return
+
+  var newLocation = agent.location
+  let lastAction = agent.cfg.getLastAction(visible)
+  if lastAction == agent.cfg.actions.moveNorth:
+    newLocation.y -= 1
+  elif lastAction == agent.cfg.actions.moveSouth:
+    newLocation.y += 1
+  elif lastAction == agent.cfg.actions.moveWest:
+    newLocation.x -= 1
+  elif lastAction == agent.cfg.actions.moveEast:
+    newLocation.x += 1
+
+  block bumpCheck:
+    agent.bump = false
+    for x in -5 .. 5:
+      for y in -5 .. 5:
+        let visibleLocation = Location(x: x, y: y)
+        let mapLocation = Location(x: x + newLocation.x, y: y + newLocation.y)
+        if mapLocation notin agent.seen:
+          continue
+        var visibleTag = agent.cfg.getTag(visible, visibleLocation)
+        if visibleTag == agent.cfg.tags.agent:
+          visibleTag = -1
+        var mapTag = agent.cfg.getTag(agent.map, mapLocation)
+        if mapTag == agent.cfg.tags.agent:
+          mapTag = -1
+        if visibleTag != mapTag:
+          newLocation = agent.location
+          agent.bump = true
+          break bumpCheck
+
+  agent.location = newLocation
+  for x in -5 .. 5:
+    for y in -5 .. 5:
+      let visibleLocation = Location(x: x, y: y)
+      let mapLocation = Location(x: x + agent.location.x, y: y + agent.location.y)
+      agent.map[mapLocation] = visible.getOrDefault(visibleLocation, @[])
+      agent.seen.incl(mapLocation)
+
+proc moveTo(agent: CogsguardAlignAllAgent, target: Location): int =
+  if agent.location == target:
+    return agent.cfg.actions.noop
+  let action = agent.cfg.aStar(agent.location, target, agent.map)
+  if action.isSome():
+    return action.get()
+  if DebugEnabled and (agent.agentId == DebugScramblerId or agent.agentId == DebugAlignerId or agent.agentId == DebugMinerId) and agent.stepCount mod DebugEvery == 0:
+    echo "[align-all] astar failed id=", agent.agentId,
+      " loc=(", agent.location.x, ",", agent.location.y, ")",
+      " target=(", target.x, ",", target.y, ")"
+  return agent.cfg.actions.noop
+
+proc stepAction(agent: CogsguardAlignAllAgent, fromLoc, toLoc: Location): int =
+  if toLoc.x == fromLoc.x + 1 and toLoc.y == fromLoc.y:
+    return agent.cfg.actions.moveEast
+  if toLoc.x == fromLoc.x - 1 and toLoc.y == fromLoc.y:
+    return agent.cfg.actions.moveWest
+  if toLoc.y == fromLoc.y - 1 and toLoc.x == fromLoc.x:
+    return agent.cfg.actions.moveNorth
+  if toLoc.y == fromLoc.y + 1 and toLoc.x == fromLoc.x:
+    return agent.cfg.actions.moveSouth
+  return agent.cfg.actions.noop
+
+proc explore(agent: CogsguardAlignAllAgent): int =
+  let unseen = agent.cfg.getNearbyUnseen(agent.location, agent.map, agent.seen, agent.unreachables)
+  if unseen.isSome():
+    let action = agent.moveTo(unseen.get())
+    if action != agent.cfg.actions.noop:
+      return action
+    agent.unreachables.incl(unseen.get())
+
+  if agent.exploreSteps < ExploreSteps:
+    let offset = Offsets4[agent.exploreDirIndex]
+    let nextLoc = agent.location + offset
+    if agent.cfg.isWalkable(agent.map, nextLoc):
+      agent.exploreSteps += 1
+      return agent.stepAction(agent.location, nextLoc)
+
+  for i in 1 .. 4:
+    let idx = (agent.exploreDirIndex + i) mod 4
+    let offset = Offsets4[idx]
+    let nextLoc = agent.location + offset
+    if agent.cfg.isWalkable(agent.map, nextLoc):
+      agent.exploreDirIndex = idx
+      agent.exploreSteps = 1
+      return agent.stepAction(agent.location, nextLoc)
+
+  return agent.cfg.actions.noop
+
+proc nearestLocation(
+  agent: CogsguardAlignAllAgent,
+  locations: seq[Location]
+): Option[Location] =
+  var bestDist = high(int)
+  var best: Option[Location] = none(Location)
+  for loc in locations:
+    let dist = manhattan(agent.location, loc)
+    if dist < bestDist:
+      bestDist = dist
+      best = some(loc)
+  return best
+
+proc nearestDepot(agent: CogsguardAlignAllAgent, alignmentFilter: int): Option[Location] =
+  var candidates: seq[Location] = @[]
+  for loc, alignment in agent.depots:
+    if alignment == alignmentFilter:
+      candidates.add(loc)
+  if candidates.len == 0:
+    return none(Location)
+  return agent.nearestLocation(candidates)
+
+proc moveToOrExplore(agent: CogsguardAlignAllAgent, target: Option[Location]): int =
+  if target.isSome():
+    return agent.moveTo(target.get())
+  return agent.explore()
+
+proc returnToHubSearch(agent: CogsguardAlignAllAgent): int =
+  if agent.hub.isSome():
+    let hubLoc = agent.hub.get()
+    if manhattan(agent.location, hubLoc) > 2:
+      return agent.moveTo(hubLoc)
+  return agent.explore()
+
+proc nearestAlignedDepot(agent: CogsguardAlignAllAgent): Option[Location] =
+  if agent.hub.isSome():
+    return agent.hub
+  return agent.nearestDepot(1)
+
+proc doDeposit(agent: CogsguardAlignAllAgent): int =
+  let depot = agent.nearestAlignedDepot()
+  return agent.moveToOrExplore(depot)
+
+proc doGather(agent: CogsguardAlignAllAgent): int =
+  var candidates: seq[Location] = @[]
+  for resource in ResourceNames:
+    for loc in agent.extractors.getOrDefault(resource, @[]):
+      let remaining = agent.extractorRemaining.getOrDefault(loc, 1)
+      if remaining != 0:
+        candidates.add(loc)
+
+  if candidates.len == 0:
+    return agent.explore()
+
+  let target = agent.nearestLocation(candidates)
+  if target.isSome():
+    return agent.moveTo(target.get())
+  return agent.explore()
+
+proc actMiner(
+  agent: CogsguardAlignAllAgent,
+  cargo: int,
+  invMiner: int
+): int =
+  if invMiner == 0:
+    if "miner_station" in agent.stations:
+      return agent.moveToOrExplore(some(agent.stations["miner_station"]))
+    return agent.returnToHubSearch()
+
+  let capacity = max(4, 40 * invMiner)
+  if cargo >= capacity - 2:
+    return agent.doDeposit()
+  return agent.doGather()
+
+proc roleForAgent(agentId: int): string =
+  case agentId
+  of 0:
+    return "scrambler"
+  of 1:
+    return "aligner"
+  of 2, 3, 4:
+    return "miner"
+  else:
+    return "scout"
+
+proc maybeMarkPending(
+  agent: CogsguardAlignAllAgent,
+  target: Location,
+  actionKind: PendingAction,
+  invHeart: int,
+  action: int
+) =
+  if action == agent.cfg.actions.noop:
+    return
+  if manhattan(agent.location, target) == 1:
+    agent.pendingAction = actionKind
+    agent.pendingTarget = some(target)
+    agent.lastHeart = invHeart
+
+proc attemptAlign(
+  agent: CogsguardAlignAllAgent,
+  target: Location,
+  invHeart: int
+): int =
+  if agent.location == target:
+    return agent.explore()
+  let action = agent.moveTo(target)
+  agent.maybeMarkPending(target, PendingAlign, invHeart, action)
+  return action
+
+proc attemptScramble(
+  agent: CogsguardAlignAllAgent,
+  target: Location,
+  invHeart: int
+): int =
+  if agent.location == target:
+    return agent.explore()
+  let action = agent.moveTo(target)
+  agent.maybeMarkPending(target, PendingScramble, invHeart, action)
+  return action
+
+proc parseVisible(
+  numTokens: int,
+  sizeToken: int,
+  rawObservation: pointer
+): Table[Location, seq[FeatureValue]] =
+  let observations = cast[ptr UncheckedArray[uint8]](rawObservation)
+  for token in 0 ..< numTokens:
+    let locationPacked = observations[token * sizeToken]
+    let featureId = observations[token * sizeToken + 1]
+    let value = observations[token * sizeToken + 2]
+    if locationPacked == 255 and featureId == 255 and value == 255:
+      break
+    var location: Location
+    if locationPacked != 0xFF:
+      location.y = (locationPacked shr 4).int - 5
+      location.x = (locationPacked and 0x0F).int - 5
+    if location notin result:
+      result[location] = @[]
+    result[location].add(FeatureValue(featureId: featureId.int, value: value.int))
+
+proc step*(
+  agent: CogsguardAlignAllAgent,
+  numAgents: int,
+  numTokens: int,
+  sizeToken: int,
+  rawObservation: pointer,
+  numActions: int,
+  agentAction: ptr int32
+) {.measure.} =
+  try:
+    discard numAgents
+    discard numActions
+
+    agent.stepCount += 1
+    let visible = parseVisible(numTokens, sizeToken, rawObservation)
+    agent.updateMap(visible)
+    agent.updateDiscoveries(visible)
+
+    let vibeId = agent.cfg.getVibe(visible, Location(x: 0, y: 0))
+    let vibeName = agent.getVibeName(vibeId)
+
+    let invHeart = agent.cfg.getInventory(visible, agent.cfg.features.invHeart)
+    let invInfluence = agent.cfg.getInventory(visible, agent.cfg.features.invInfluence)
+    let invAligner = agent.cfg.getInventory(visible, agent.cfg.features.invAligner)
+    let invScrambler = agent.cfg.getInventory(visible, agent.cfg.features.invScrambler)
+    let invMiner = agent.cfg.getInventory(visible, agent.cfg.features.invMiner)
+    let invCarbon = agent.cfg.getInventory(visible, agent.cfg.features.invCarbon)
+    let invOxygen = agent.cfg.getInventory(visible, agent.cfg.features.invOxygen)
+    let invGermanium = agent.cfg.getInventory(visible, agent.cfg.features.invGermanium)
+    let invSilicon = agent.cfg.getInventory(visible, agent.cfg.features.invSilicon)
+    let cargo = invCarbon + invOxygen + invGermanium + invSilicon
+
+    if agent.pendingAction != PendingNone and invHeart < agent.lastHeart:
+      if agent.pendingTarget.isSome():
+        let target = agent.pendingTarget.get()
+        if agent.pendingAction == PendingAlign:
+          agent.depots[target] = 1
+        elif agent.pendingAction == PendingScramble:
+          agent.depots[target] = 0
+      agent.pendingAction = PendingNone
+      agent.pendingTarget = none(Location)
+
+    agent.lastHeart = invHeart
+
+    var action = agent.cfg.actions.noop
+
+    let targetScramble = agent.nearestDepot(-1)
+    let targetAlign = agent.nearestDepot(0)
+    let role = roleForAgent(agent.agentId)
+
+    if DebugEnabled and (agent.agentId == DebugScramblerId or agent.agentId == DebugAlignerId or agent.agentId == DebugMinerId) and agent.stepCount mod DebugEvery == 0:
+      var clipsDepots = 0
+      var neutralDepots = 0
+      var cogsDepots = 0
+      for _, alignment in agent.depots:
+        if alignment == -1:
+          clipsDepots += 1
+        elif alignment == 0:
+          neutralDepots += 1
+        elif alignment == 1:
+          cogsDepots += 1
+      let hasScramblerStation = "scrambler_station" in agent.stations
+      let hasAlignerStation = "aligner_station" in agent.stations
+      let hasMinerStation = "miner_station" in agent.stations
+      let chestDist = if agent.chest.isSome(): manhattan(agent.location, agent.chest.get()) else: -1
+      let hubDist = if agent.hub.isSome(): manhattan(agent.location, agent.hub.get()) else: -1
+      let minerCapacity = max(4, 40 * invMiner)
+      echo "[align-all] step=", agent.stepCount,
+        " id=", agent.agentId,
+        " role=", role,
+        " vibe=", vibeName,
+        " loc=(", agent.location.x, ",", agent.location.y, ")",
+        " heart=", invHeart,
+        " infl=", invInfluence,
+        " gearA=", invAligner,
+        " gearS=", invScrambler,
+        " gearM=", invMiner,
+        " cargo=", cargo, "/", minerCapacity,
+        " depots=", clipsDepots, "/", neutralDepots, "/", cogsDepots,
+        " chest=", agent.chest.isSome(),
+        " hub=", agent.hub.isSome(),
+        " stations(a/s/m)=", hasAlignerStation, "/", hasScramblerStation, "/", hasMinerStation,
+        " dist(chest/hub)=", chestDist, "/", hubDist,
+        " targetA=", targetAlign.isSome(),
+        " targetS=", targetScramble.isSome()
+
+    if role == "scrambler":
+      if vibeName != "scrambler":
+        action = agent.actionForVibe("scrambler")
+      elif invScrambler == 0:
+        if "scrambler_station" in agent.stations:
+          action = agent.moveToOrExplore(some(agent.stations["scrambler_station"]))
+        else:
+          action = agent.returnToHubSearch()
+      elif invHeart == 0:
+        if agent.chest.isSome():
+          action = agent.moveToOrExplore(agent.chest)
+        else:
+          action = agent.returnToHubSearch()
+      elif targetScramble.isSome():
+        action = agent.attemptScramble(targetScramble.get(), invHeart)
+      else:
+        action = agent.explore()
+    elif role == "aligner":
+      if vibeName != "aligner":
+        action = agent.actionForVibe("aligner")
+      elif invAligner == 0:
+        if "aligner_station" in agent.stations:
+          action = agent.moveToOrExplore(some(agent.stations["aligner_station"]))
+        else:
+          action = agent.returnToHubSearch()
+      elif invHeart == 0:
+        if agent.chest.isSome():
+          action = agent.moveToOrExplore(agent.chest)
+        else:
+          action = agent.returnToHubSearch()
+      elif invInfluence == 0:
+        if agent.hub.isSome():
+          action = agent.moveToOrExplore(agent.hub)
+        else:
+          action = agent.returnToHubSearch()
+      elif targetAlign.isSome():
+        action = agent.attemptAlign(targetAlign.get(), invHeart)
+      else:
+        action = agent.explore()
+    elif role == "miner":
+      if vibeName != "miner":
+        action = agent.actionForVibe("miner")
+      else:
+        action = agent.actMiner(cargo, invMiner)
+    else:
+      if vibeName != "scout":
+        action = agent.actionForVibe("scout")
+      else:
+        action = agent.explore()
+
+    agentAction[] = action.int32
+  except:
+    echo getCurrentException().getStackTrace()
+    echo getCurrentExceptionMsg()
+    quit()
+
+proc newCogsguardAlignAllAgent*(agentId: int, environmentConfig: string): CogsguardAlignAllAgent =
+  var config = parseConfig(environmentConfig)
+  result = CogsguardAlignAllAgent(agentId: agentId, cfg: config)
+  result.random = initRand(agentId)
+  result.map = initTable[Location, seq[FeatureValue]]()
+  result.seen = initHashSet[Location]()
+  result.unreachables = initHashSet[Location]()
+  result.location = Location(x: 0, y: 0)
+  result.exploreDirIndex = 0
+  result.exploreSteps = 0
+  result.stations = initTable[string, Location]()
+  result.depots = initTable[Location, int]()
+  result.extractors = initTable[string, seq[Location]]()
+  result.extractorRemaining = initTable[Location, int]()
+  result.hub = none(Location)
+  result.chest = none(Location)
+  result.actionIds = initTable[string, int]()
+  for id, name in config.config.actions:
+    result.actionIds[name] = id
+  result.pendingAction = PendingNone
+  result.pendingTarget = none(Location)
+  result.lastHeart = 0
+  result.stepCount = 0
+
+proc newCogsguardAlignAllPolicy*(environmentConfig: string): CogsguardAlignAllPolicy =
+  let cfg = parseConfig(environmentConfig)
+  var agents: seq[CogsguardAlignAllAgent] = @[]
+  for id in 0 ..< cfg.config.numAgents:
+    agents.add(newCogsguardAlignAllAgent(id, environmentConfig))
+  CogsguardAlignAllPolicy(agents: agents)
+
+proc stepBatch*(
+  policy: CogsguardAlignAllPolicy,
+  agentIds: pointer,
+  numAgentIds: int,
+  numAgents: int,
+  numTokens: int,
+  sizeToken: int,
+  rawObservations: pointer,
+  numActions: int,
+  rawActions: pointer
+) =
+  let ids = cast[ptr UncheckedArray[int32]](agentIds)
+  let obsArray = cast[ptr UncheckedArray[uint8]](rawObservations)
+  let actionArray = cast[ptr UncheckedArray[int32]](rawActions)
+  let obsStride = numTokens * sizeToken
+
+  for i in 0 ..< numAgentIds:
+    let idx = int(ids[i])
+    let obsPtr = cast[pointer](obsArray[idx * obsStride].addr)
+    let actPtr = cast[ptr int32](actionArray[idx].addr)
+    step(policy.agents[idx], numAgents, numTokens, sizeToken, obsPtr, numActions, actPtr)