OSCR

Decoding behavior with minimal and interpretable agent models.

Code ↔ Paper

6 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 6 matches
  1. [1] § Methods › Metric-adaptive particle swarm optimization ↔ src/DiscreteObs/inference.py, lines 287–346 · score 0.79 · social coefficient, cognitive coefficient, inertia, swarming, global, neighbors
  2. [2] § Methods › Metric-adaptive particle swarm optimization ↔ src/MAPSO.py, lines 341–418 · score 0.65 · nearest, coefficient, inertia, social, swarming, cognitive
  3. [3] § Methods › Metric-adaptive particle swarm optimization ↔ src/DiscreteObs/inference.py, lines 287–346 · score 0.63 · global best, parameter space, mutation, adaptive, swarm, neighbors
  4. [4] § Methods › MAPSO training schedule ↔ src/FSC.py, lines 852–915 · score 0.55 · best inferred, negative log likelihood, MAPSO, epochs, optimization, particle
  5. [5] § Methods › Metric-adaptive particle swarm optimization ↔ src/MAPSO.py, lines 119–187 · score 0.55 · mutated, mutation, adaptively, swarming, MAPSO, velocity
  6. [6] § Methods › MAPSO training schedule ↔ src/FSC.py, lines 713–771 · score 0.52 · Gaussian distribution, particles initial, multivariate, covariance, MAPSO

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 929 lines · 41 KB · no license · 2 matches

  1. import torch
  2. from torch import nn
  3. import numpy as np
  4. import numba as nb
  5. import random
  7. import utils
  8. import MAPSO
  10. class InferenceDiscreteObs:
  11. def __init__(self, FSC):
  12. """
  13. Initialize the inference backend for a discrete-observation FSC.
  15. This constructor moves all FSC trainable quantities to the best
  16. available torch device (CUDA, MPS, or CPU), converts numpy arrays to
  17. tensors when needed, and wraps parameters in ``nn.Parameter`` if the
  18. FSC is not already trained.
  20. Parameters:
  21. --- FSC: FSC
  22. Parent FSC object configured for discrete observations.
  23. """
  24. self.FSC = FSC
  26. if torch.cuda.is_available():
  27. self.device = torch.device("cuda")
  28. elif torch.backends.mps.is_available():
  29. self.device = torch.device("mps")
  30. else:
  31. self.device = torch.device("cpu")
  34. if not isinstance(self.FSC.psi, torch.Tensor):
  35. self.FSC.psi = torch.tensor(self.FSC.psi.astype(np.float32), device=self.device)
  37. if not isinstance(self.FSC.psi, nn.Parameter) and FSC.trained == False:
  38. self.FSC.psi = nn.Parameter(self.FSC.psi)
  40. for idx, param in enumerate(self.FSC.GPModel.params):
  41. param_name = self.FSC.GPModel.param_names[idx]
  43. if not isinstance(param, torch.Tensor):
  44. param = torch.tensor(param.astype(np.float32), device=self.device)
  46. if not isinstance(param, nn.Parameter) and FSC.trained == False:
  47. param = nn.Parameter(param)
  49. self.FSC.GPModel.__setattr__(param_name, param)
  51. self.InternalMemSpace = torch.arange(self.FSC.M)
  52. self.InternalActSpace = torch.arange(self.FSC.A)
  53. self.InternalObsSpace = torch.arange(self.FSC.Y)
  55. self.trajectories_loaded = False
  56. self.optimizer_initialized = False
  57. self.trained = False
  59. def get_policy_params(self):
  60. """
  61. Return the current policy parameters in model-defined order.
  63. Returns:
  64. --- tuple of torch.Tensor
  65. Tuple ``(param_0, ..., param_k)`` matching
  66. ``self.FSC.GPModel.param_names``.
  67. """
  68. return tuple([self.FSC.GPModel.__getattribute__(param) for param in self.FSC.GPModel.param_names])
  70. def get_TMat(self):
  71. """
  72. Return the joint transition matrix used during inference.
  74. Returns:
  75. --- torch.Tensor
  76. Transition tensor with shape ``(Y, A, M, M)``.
  77. """
  78. return self.FSC.GPModel.get_TMat_torch()
  80. def get_memory_transition(self):
  81. """
  82. Return the memory transition component of the policy.
  84. Returns:
  85. --- torch.Tensor
  86. Memory transition tensor ``g(m' | a, m, y)`` with shape
  87. ``(Y, A, M, M)``.
  88. """
  89. return self.FSC.GPModel.get_memory_transition_torch()
  91. def get_action_policy(self):
  92. """
  93. Return the marginal action policy.
  95. Returns:
  96. --- torch.Tensor
  97. Action policy tensor ``pi(a | m)`` with shape ``(M, A)``.
  98. """
  99. return self.FSC.GPModel.get_action_policy_torch()
  102. def load_trajectories(self, trajectories):
  103. """
  104. Loads a set of trajectories to be used for training the FSC.
  106. Parameters:
  107. --- trajectories: list of dicts
  108. List of dictionaries containing the actions and observations for each trajectory.
  109. """
  110. self.ObsAct_trajectories = []
  111. self.observations_trajectories_np = []
  112. self.actions_trajectories_np = []
  113. self.n_trajectories = len(trajectories)
  115. self.pStart_ya_emp = np.zeros((self.FSC.Y, self.FSC.A))
  118. for trajectory in trajectories:
  119. observations = self._map_obs_to_internal_space(trajectory["observations"])
  120. actions = self._map_act_to_internal_space(trajectory["actions"])
  121. self.observations_trajectories_np.append(observations)
  122. self.actions_trajectories_np.append(actions)
  124. self.ObsAct_trajectories.append([torch.tensor(observations), torch.tensor(actions)])
  125. y0 = observations[0]
  126. a0 = actions[0]
  128. self.pStart_ya_emp[y0, a0] += 1
  130. self.pStart_ya_emp /= np.sum(self.pStart_ya_emp)
  132. self.trajectories_loaded = True
  134. def _map_obs_to_internal_space(self, obs):
  135. """
  136. Helper method to map an observation sequence to the internal observation space.
  138. Parameters:
  139. --- obs: np.array
  140. Observation sequence to map.
  142. Returns:
  143. --- obs_internal: np.array
  144. Observation sequence in the internal observation space.
  145. """
  146. return np.array([self.get_obs_idx(o) for o in obs])
  148. def _map_act_to_internal_space(self, act):
  149. """
  150. Helper method to map an action sequence to the internal action space.
  152. Parameters:
  153. --- act: np.array
  154. Action sequence to map.
  156. Returns:
  157. --- act_internal: np.array
  158. Action sequence in the internal action space.
  159. """
  160. return np.array([self.get_act_idx(a) for a in act])
  162. def get_obs_idx(self, obs):
  163. """
  164. Helper method to get the index of an observation in the observation space.
  166. Parameters:
  167. --- obs: int
  168. Observation to get the index of.
  170. Returns:
  171. --- idx: int
  172. Index of the observation in the observation space.
  173. """
  175. return np.where(obs == self.FSC.ObsSpace)[0][0]
  177. def get_act_idx(self, act):
  178. """
  179. Helper method to get the index of an action in the action space.
  181. Parameters:
  182. --- act: int
  183. Action to get the index of.
  185. Returns:
  186. --- idx: int
  187. Index of the action in the action space.
  188. """
  190. return np.where(self.FSC.ActSpace == act)[0][0]
  192. def get_mem_idx(self, mem):
  193. """
  194. Helper method to get the index of a memory state in the memory space.
  196. Parameters:
  197. --- mem: int
  198. Memory state to get the index of.
  200. Returns:
  201. --- idx: int
  202. Index of the memory state in the memory space
  203. """
  205. return np.where(self.FSC.MemSpace == mem)[0][0]
  207. def evaluate_nloglikelihood(self, idx_traj, grad_required=False):
  208. """
  209. Wrapper method to evaluate the negative log-likelihood of a given trajectory.
  210. It distinguishes between the case of custom observation and action spaces and the case
  211. of default observation and action spaces.
  213. Parameters:
  214. --- idx_traj: int
  215. Index of the trajectory to evaluate.
  216. --- grad_required: bool (default = False)
  217. Flag indicating whether the gradient is required or not.
  219. Returns:
  220. --- nLL: float
  221. Negative log-likelihood of the trajectory.
  222. """
  223. observations, actions = self.ObsAct_trajectories[idx_traj]
  225. return self.loss(observations, actions, grad_required = grad_required)
  227. def loss(self, observations, actions, grad_required=True):
  228. """
  229. Method to compute the negative log-likelihood of a given trajectory with default observation and action spaces.
  230. The gradients of the loss are computed if the grad_required flag is set to True.
  232. Parameters:
  233. --- observations: torch.tensor
  234. Array of observations.
  235. --- actions: torch.tensor
  236. Array of actions.
  237. --- grad_required: bool (default = True)
  238. Flag indicating whether the gradient is required or not.
  240. Returns:
  241. --- nLL: float
  242. Negative log-likelihood of the trajectory.
  243. """
  244. nLL = torch.tensor(0.0, requires_grad = grad_required)
  246. TMat = self.FSC.GPModel.get_TMat_torch()
  248. if self.FSC._init_memory_obs_dependent:
  249. rho = self.FSC.rho[observations[0]]
  250. else:
  251. rho = self.FSC.rho
  253. for t in range(observations.size(0)):
  254. idx_a = actions[t]
  255. idx_obs = observations[t]
  257. transition_probs = TMat[idx_obs, idx_a].T
  259. if torch.sum(transition_probs) == 0:
  260. break
  262. if t == 0:
  263. if transition_probs.device.type == 'mps':
  264. # MPS-specific workaround
  265. transition_probs_safe = transition_probs.clone().detach().requires_grad_(transition_probs.requires_grad)
  266. rho_safe = rho.clone().detach().requires_grad_(rho.requires_grad)
  267. m = torch.matmul(transition_probs_safe, rho_safe)
  268. else:
  269. m = torch.matmul(transition_probs, rho)
  270. else:
  271. m = torch.matmul(transition_probs, m)
  273. if torch.sum(m) == 0:
  274. break
  276. mv = torch.sum(m)
  277. nLL = nLL - torch.log(mv)
  278. m /= mv
  281. if torch.sum(m) == 0:
  282. return nLL
  283. else:
  284. return nLL - torch.log(torch.sum(m))
  287. def optimize_w_MAPSO(self, trainable_params, trainable_params_mask,
  288. n_particles, NEpochs,
  289. init_particles, init_velocities,
  290. c1_init, c2_init, w_init,
  291. sigma_min, sigma_max,
  292. dynamic_topology, n_neighbors_init, n_neighbors_final, num_neighbors_mid,
  293. verbose, verbose_epochs):
  294. """
  295. Optimize FSC parameters with MAPSO (Adaptive Particle Swarm).
  297. This method builds a flattened parameter space containing all policy
  298. parameters plus ``psi``, applies optional element-wise trainability
  299. masks, initializes particle positions/velocities from the configured
  300. distributions, and runs either global-best MAPSO or kNN-local MAPSO.
  302. Parameters:
  303. --- trainable_params: dict
  304. Dictionary mapping parameter names (including ``psi``) to booleans
  305. indicating whether each parameter block is trainable.
  306. --- trainable_params_mask: dict
  307. Optional element-wise masks. Values are either ``None`` or boolean
  308. arrays with the same shape as the corresponding parameter.
  309. --- n_particles: int
  310. Number of particles in the swarm.
  311. --- NEpochs: int
  312. Number of MAPSO iterations.
  313. --- init_particles: dict
  314. Particle initialization settings (distribution and hyperparameters).
  315. --- init_velocities: dict
  316. Velocity initialization settings (distribution and hyperparameters).
  317. --- c1_init: float
  318. Initial cognitive coefficient.
  319. --- c2_init: float
  320. Initial social coefficient.
  321. --- w_init: float
  322. Initial inertia weight.
  323. --- sigma_min: float
  324. Minimum mutation scale used by MAPSO convergence strategy.
  325. --- sigma_max: float
  326. Maximum mutation scale used by MAPSO convergence strategy.
  327. --- dynamic_topology: bool
  328. If True, use kNN local-best topology; otherwise use global-best.
  329. --- n_neighbors_init: int or None
  330. Initial number of neighbors for dynamic topology.
  331. --- n_neighbors_final: int or None
  332. Final number of neighbors for dynamic topology.
  333. --- num_neighbors_mid: int or None
  334. Midpoint number of neighbors for dynamic topology.
  335. --- verbose: bool
  336. Print per-iteration diagnostics.
  337. --- verbose_epochs: bool
  338. Print per-epoch summary.
  340. Returns:
  341. --- np.ndarray
  342. Best objective value at each MAPSO iteration.
  343. """
  345. assert self.trajectories_loaded, "No trajectories have been loaded. Load trajectories with the load_trajectories method."
  346. assert not self.trained, "The model has already been trained. If you want to train it again, reinitialize it or set the flag self.trained to False."
  348. spacedim = sum([param.numel() for param in self.get_policy_params()]) + self.FSC.psi.numel()
  350. trainable_mask = np.zeros(spacedim, dtype=bool)
  351. init_pos = np.zeros((n_particles, spacedim))
  352. init_vel = np.zeros((n_particles, spacedim))
  354. if init_particles["distribution"] == "multivariate_normal":
  355. random_pos_mv = np.random.multivariate_normal(init_particles["mean"], init_particles["cov"],
  356. n_particles)
  358. start_idx = 0
  359. for idx, param in enumerate(self.get_policy_params()):
  360. param_name = self.FSC.GPModel.param_names[idx]
  361. end_idx = start_idx + param.numel()
  362. if trainable_params_mask[param_name] is not None:
  363. trainable_mask[start_idx:end_idx] = trainable_params_mask[param_name].flatten()
  364. else:
  365. trainable_mask[start_idx:end_idx] = trainable_params[param_name]
  367. flatten_param = self.FSC.GPModel.__getattribute__(self.FSC.GPModel.param_names[idx]).detach().cpu().numpy().flatten()
  368. init_pos[:, start_idx:end_idx] = np.tile(flatten_param, (n_particles, 1))
  369. init_vel[:, start_idx:end_idx] = np.zeros((n_particles, param.numel()))
  371. if trainable_params[param_name]:
  372. num_trainable = np.sum(trainable_mask[start_idx:end_idx])
  374. if init_particles["distribution"] == "uniform":
  375. random_pos = np.random.uniform(init_particles["xmin"], init_particles["xmax"],
  376. (n_particles, num_trainable))
  377. elif init_particles["distribution"] == "normal":
  378. random_pos = np.random.normal(init_particles["mean"], init_particles["std"],
  379. (n_particles, num_trainable))
  380. elif init_particles["distribution"] == "multivariate_normal":
  381. random_pos = random_pos_mv[:, start_idx:end_idx][:, trainable_mask[start_idx:end_idx]]
  382. elif init_particles["distribution"] == "uniform_with_biases":
  383. random_pos = np.random.uniform(init_particles["xmin"], init_particles["xmax"],
  384. (n_particles, num_trainable))
  385. random_pos += init_particles["biases"][start_idx:end_idx][trainable_mask[start_idx:end_idx]]
  386. else:
  387. raise ValueError("Invalid position distribution.")
  389. if init_velocities["distribution"] == "uniform":
  390. random_vel = np.random.uniform(init_velocities["vmin"], init_velocities["vmax"],
  391. (n_particles, num_trainable))
  392. elif init_velocities["distribution"] == "normal":
  393. random_vel = np.random.normal(init_velocities["mean"], init_velocities["std"],
  394. (n_particles, num_trainable))
  395. else:
  396. raise ValueError("Invalid velocity distribution.")
  398. init_pos[:, start_idx:end_idx][:, trainable_mask[start_idx:end_idx]] = random_pos
  399. init_vel[:, start_idx:end_idx][:, trainable_mask[start_idx:end_idx]] = random_vel
  401. start_idx = end_idx
  403. if trainable_params_mask["psi"] is not None:
  404. trainable_mask[-self.FSC.psi.numel():] = trainable_params_mask["psi"].flatten()
  405. else:
  406. trainable_mask[-self.FSC.psi.numel():] = trainable_params["psi"]
  408. #print(self.FSC.psi.shape, self.FSC.psi.numel())
  409. init_pos[:, -self.FSC.psi.numel():] = np.tile(self.FSC.psi.detach().cpu().numpy().flatten(), (n_particles, 1))
  410. init_vel[:, -self.FSC.psi.numel():] = np.zeros((n_particles, self.FSC.psi.numel()))
  412. if trainable_params["psi"]:
  413. num_trainable = np.sum(trainable_mask[-self.FSC.psi.numel():])
  415. if init_particles["distribution"] == "uniform":
  416. random_pos = np.random.uniform(init_particles["xmin"], init_particles["xmax"],
  417. (n_particles, num_trainable))
  418. elif init_particles["distribution"] == "normal":
  419. random_pos = np.random.normal(init_particles["mean"], init_particles["std"],
  420. (n_particles, num_trainable))
  421. elif init_particles["distribution"] == "multivariate_normal":
  422. random_pos = random_pos_mv[:, -self.FSC.psi.numel():][:, trainable_mask[-self.FSC.psi.numel():]]
  423. elif init_particles["distribution"] == "uniform_with_biases":
  424. random_pos = np.random.uniform(init_particles["xmin"], init_particles["xmax"],
  425. (n_particles, num_trainable))
  426. random_pos += init_particles["biases"][-self.FSC.psi.numel():][trainable_mask[-self.FSC.psi.numel():]]
  427. else:
  428. raise ValueError("Invalid position distribution.")
  430. if init_velocities["distribution"] == "uniform":
  431. random_vel = np.random.uniform(init_velocities["vmin"], init_velocities["vmax"],
  432. (n_particles, num_trainable))
  433. elif init_velocities["distribution"] == "normal":
  434. random_vel = np.random.normal(init_velocities["mean"], init_velocities["std"],
  435. (n_particles, num_trainable))
  436. else:
  437. raise ValueError("Invalid velocity distribution.")
  439. init_pos[:, -self.FSC.psi.numel():][:, trainable_mask[-self.FSC.psi.numel():]] = random_pos
  440. init_vel[:, -self.FSC.psi.numel():][:, trainable_mask[-self.FSC.psi.numel():]] = random_vel
  443. if dynamic_topology:
  444. gbests, gbest_values = MAPSO.particle_swarm_optimization_discrete_kNN(self.FSC.GPModel._nb_get_TMat_flat,
  445. trainable_mask,
  446. spacedim, n_particles, NEpochs,
  447. self.FSC.M, self.FSC.A, self.FSC.Y,
  448. self.observations_trajectories_np,
  449. self.actions_trajectories_np,
  450. init_pos, init_vel,
  451. num_neighbors_init = n_neighbors_init,
  452. num_neighbors_final = n_neighbors_final,
  453. num_neighbors_mid = num_neighbors_mid,
  454. c1 = c1_init, c2 = c2_init, w = w_init,
  455. sigma_min = sigma_min, sigma_max = sigma_max,
  456. verbose = verbose, verbose_epochs = verbose_epochs,
  457. init_memory_obs_dependent = self.FSC._init_memory_obs_dependent)
  458. else:
  459. gbests, gbest_values = MAPSO.particle_swarm_optimization_discrete(self.FSC.GPModel._nb_get_TMat_flat,
  460. trainable_mask,
  461. spacedim, n_particles, NEpochs,
  462. self.FSC.M, self.FSC.A, self.FSC.Y,
  463. self.observations_trajectories_np,
  464. self.actions_trajectories_np,
  465. init_pos, init_vel,
  466. c1 = c1_init, c2 = c2_init, w = w_init,
  467. sigma_min = sigma_min, sigma_max = sigma_max,
  468. verbose = verbose, verbose_epochs = verbose_epochs,
  469. init_memory_obs_dependent = self.FSC._init_memory_obs_dependent)
  471. for idx, param in enumerate(self.get_policy_params()):
  472. param_name = self.FSC.GPModel.param_names[idx]
  473. start_idx = 0
  474. for idx, param in enumerate(self.get_policy_params()):
  475. param_name = self.FSC.GPModel.param_names[idx]
  476. end_idx = start_idx + param.numel()
  477. new_param = torch.tensor(gbests[-1, start_idx:end_idx].reshape(param.shape).astype(np.float32), device=self.device)
  478. self.FSC.GPModel.__setattr__(param_name,
  479. nn.Parameter(new_param))
  480. start_idx = end_idx
  481. new_psi = gbests[-1, -self.FSC.psi.numel():].astype(np.float32)
  482. new_psi = new_psi.reshape(self.FSC.psi.shape)
  483. self.FSC.psi = nn.Parameter(torch.tensor(new_psi, device=self.device))
  485. return gbest_values
  487. def optimize_w_gradient(self, use_ccopt,
  488. trainable_params, trainable_params_mask, any_masked,
  489. NEpochs, NBatch, lr,
  490. train_split, optimizer, scheduler_dict,
  491. maxiter, rho0, th, c_gauge,
  492. verbose, verbose_epochs):
  493. """
  494. Optimize FSC parameters with gradient-based training.
  496. Supports per-parameter learning rates, optional parameter masking, an
  497. optional convex-concave optimization step for ``rho`` (when enabled),
  498. mini-batch training, and optional validation split with best-epoch
  499. model selection.
  501. Parameters:
  502. --- use_ccopt: bool
  503. If True, update ``psi`` through convex-concave optimization of rho
  504. at each batch (only for observation-independent initialization).
  505. --- trainable_params: dict
  506. Parameter-level trainability flags.
  507. --- trainable_params_mask: dict
  508. Optional element-wise trainability masks.
  509. --- any_masked: bool
  510. True if at least one element-wise mask is active.
  511. --- NEpochs: int
  512. Number of gradient epochs.
  513. --- NBatch: int
  514. Batch size in number of trajectories.
  515. --- lr: float or dict
  516. Global learning rate or per-parameter learning-rate dictionary.
  517. --- train_split: float
  518. Fraction of trajectories used for training.
  519. --- optimizer: str
  520. Optimizer name (expected: ``ADAM`` or ``SGD``).
  521. --- scheduler_dict: dict
  522. Scheduler configuration dictionary.
  523. --- maxiter: int or None
  524. Maximum iterations for ccopt.
  525. --- rho0: np.ndarray or None
  526. Initial rho for ccopt.
  527. --- th: float or None
  528. Convergence threshold for ccopt.
  529. --- c_gauge: float or None
  530. Additive gauge applied to ``log(rho)`` when reconstructing psi.
  531. --- verbose: bool
  532. Verbosity flag (kept for interface consistency).
  533. --- verbose_epochs: bool
  534. If True, print per-epoch losses and learning rate.
  536. Returns:
  537. --- tuple
  538. ``(losses_train, losses_val)`` where validation losses are ``None``
  539. when no validation split is used.
  540. """
  541. assert self.trajectories_loaded, "No trajectories have been loaded. Load trajectories with the load_trajectories method."
  542. assert not self.trained, "The model has already been trained. If you want to train it again, reinitialize it or set the flag self.trained to False."
  544. lr_dict = {}
  546. if isinstance(lr, float):
  547. single_lr = True
  548. for param in self.FSC.GPModel.param_names:
  549. lr_dict[param] = lr
  550. lr_dict["psi"] = lr
  551. elif isinstance(lr, dict):
  552. for param in self.FSC.GPModel.param_names:
  553. if param != "psi" and trainable_params[param]:
  554. if param in lr:
  555. lr_dict[param] = lr[param]
  556. else:
  557. raise ValueError(f"Missing learning rate for parameter {param}.")
  558. if "psi" not in lr and trainable_params["psi"] and not use_ccopt:
  559. raise ValueError("Missing learning rate for psi.")
  560. else:
  561. lr_dict["psi"] = lr["psi"]
  562. else:
  563. raise ValueError("Invalid learning rate. The learning rate must be a float or a dictionary with the parameters as keys.")
  565. parkey_optimizer = []
  566. for idx, param in enumerate(self.FSC.GPModel.param_names):
  567. if trainable_params[param]:
  568. parkey_optimizer.append({"params": self.FSC.GPModel.__getattribute__(param), 'lr': lr_dict[param]})
  569. if trainable_params["psi"] and not use_ccopt:
  570. parkey_optimizer.append({'params': self.FSC.psi, 'lr': lr_dict["psi"]})
  572. if optimizer == "ADAM":
  573. self.optimizer = torch.optim.Adam(parkey_optimizer)
  574. elif optimizer == "SDG":
  575. self.optimizer = torch.optim.SGD(parkey_optimizer)
  577. if scheduler_dict["type"] == "exponential":
  578. scheduler = torch.optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=scheduler_dict["decay_rate"])
  579. elif scheduler_dict["type"] == "fixed":
  580. scheduler = None
  581. else:
  582. raise ValueError("Invalid scheduler.")
  584. NTrain = int(train_split * len(self.ObsAct_trajectories))
  585. NVal = len(self.ObsAct_trajectories) - NTrain
  587. trjs_train = self.ObsAct_trajectories[:NTrain]
  588. best_params = [self.FSC.GPModel.__getattribute__(param) for param in self.FSC.GPModel.param_names]
  589. best_psi = self.FSC.psi
  590. best_epoch = 0
  592. losses_train = []
  594. init_loss = 0
  595. for idx_traj in range(NTrain):
  596. init_loss += self.loss(trjs_train[idx_traj][0], trjs_train[idx_traj][1], grad_required=False).item()
  597. losses_train.append(init_loss / NTrain)
  599. init_msg = f"Training with {NTrain} trajectories"
  600. if NVal != 0:
  601. trjs_val = self.FeatAct_trajectories[NTrain:]
  602. losses_val = []
  604. init_loss_val = 0
  605. for idx_traj in range(NVal):
  606. init_loss_val += self.loss(trjs_val[idx_traj][0], trjs_val[idx_traj][1], grad_required=False).item()
  607. losses_val.append(init_loss_val / NVal)
  609. init_msg += f" and validating with {NVal} trajectories."
  610. init_msg += " Initial training loss: " + str(losses_train[0]) + ". Initial validation loss: " + str(losses_val[0]) + "."
  611. else:
  612. init_msg += ". Initial loss: " + str(losses_train[0]) + "."
  614. if verbose_epochs:
  615. if single_lr:
  616. print(init_msg + f" Using a single learning rate of {lr}.")
  617. else:
  618. for idx, param in enumerate(self.FSC.GPModel.param_names):
  619. print(init_msg + f" Using learning rate {lr[param]} for {param}.")
  620. print(init_msg + f" Using learning rate {lr['psi']} for psi.")
  622. for epoch in range(NEpochs):
  623. running_loss = 0.0
  624. running_count = 0
  626. random.shuffle(trjs_train)
  628. for idx in range(0, NTrain, NBatch):
  629. if any_masked:
  630. pre_loss_params = {}
  631. for key in self.FSC.GPModel.param_names:
  632. mask = trainable_params_mask[key]
  633. if mask is not None:
  634. pre_loss_params[key] = self.FSC.GPModel.__getattribute__(key).detach().clone()
  635. if trainable_params_mask["psi"] is not None:
  636. pre_loss_psi = self.FSC.psi.detach().clone()
  638. self.optimizer.zero_grad()
  639. loss = torch.tensor(0.0, requires_grad=True)
  641. TMat = self.FSC.GPModel.get_TMat_torch()
  643. if use_ccopt and not self.FSC._init_memory_obs_dependent and trainable_params["psi"]:
  644. if rho0 is None:
  645. rho0 = np.ones(self.FSC.M)/self.FSC.M
  647. rho, _ = InferenceDiscreteObs.optimize_rho(self.FSC.Y, self.FSC.M, self.FSC.A,
  648. TMat.detach().cpu().numpy(), self.pStart_ya_emp,
  649. rho0, maxiter, th = th)
  651. rho = torch.tensor(rho.astype(np.float32), device = self.device)
  652. self.FSC.psi = nn.Parameter(torch.log(rho) + c_gauge)
  654. count = 0
  655. for idx_traj in range(idx, idx + NBatch):
  656. if idx_traj < NTrain:
  657. loss_traj = self.loss(trjs_train[idx_traj][0], trjs_train[idx_traj][1])
  658. if torch.isnan(loss_traj):
  659. continue
  660. loss = loss + loss_traj
  661. count += 1
  663. if count == 0:
  664. err_msg = "Gradient optimization failed because no valid trajectories were found in a batch. This means that the loss could not be evaluated due to forbidden transition, and that the current parameters are not compatible with some trajectories."
  665. err_msg += " Either improve initialization or choose a smaller learning rate. Overwriting with the best parameters found so far."
  667. self.FSC.psi = best_psi
  668. for idx, param in enumerate(self.FSC.GPModel.param_names):
  669. self.FSC.GPModel.__setattr__(param, nn.Parameter(best_params[idx]))
  671. raise RuntimeError(err_msg)
  673. loss.backward()
  674. self.optimizer.step()
  675. running_loss += loss.item()
  676. running_count += count
  678. if any_masked:
  679. for key in self.FSC.GPModel.param_names:
  680. mask = ~trainable_params_mask[key]
  681. if mask is not None:
  682. self.FSC.GPModel.__getattribute__(key).data[mask] = pre_loss_params[key].data[mask].clone()
  683. mask_psi = ~trainable_params_mask["psi"]
  684. if mask_psi is not None:
  685. self.FSC.psi.data[mask_psi] = pre_loss_psi.data[mask_psi].clone()
  687. running_loss = running_loss / running_count
  688. losses_train.append(running_loss)
  690. if NVal != 0:
  691. running_loss_val = 0.0
  693. for idx_traj in range(NVal):
  694. loss_val = torch.tensor(0.0, requires_grad=False)
  696. loss_traj_val = self.loss(trjs_val[idx_traj][0], trjs_val[idx_traj][1], grad_required=False)
  697. loss_val = loss_val + loss_traj_val
  699. running_loss_val += loss_val.item()
  701. running_loss_val = running_loss_val / NVal
  702. losses_val.append(running_loss_val)
  704. if running_loss_val < min(losses_val[:-1]):
  705. best_params = [self.FSC.GPModel.__getattribute__(param).detach().clone() for param in self.FSC.GPModel.param_names]
  706. best_psi = self.FSC.psi.detach().clone()
  707. best_epoch = epoch + 1
  709. if verbose_epochs:
  710. print(f"Epoch {epoch + 1} - Training loss: {round(running_loss, 5)}, Validation loss: {round(running_loss_val, 5)} - Learning rate: {round(self.optimizer.param_groups[0]['lr'], 5)}")
  712. else:
  713. if running_loss < min(losses_train[:-1]):
  714. best_params = [self.FSC.GPModel.__getattribute__(param).detach().clone() for param in self.FSC.GPModel.param_names]
  715. best_psi = self.FSC.psi.detach().clone()
  716. best_epoch = epoch + 1
  718. if verbose_epochs:
  719. print(f"Epoch {epoch + 1} - Training loss: {round(running_loss, 5)} - Learning rate: {round(self.optimizer.param_groups[0]['lr'], 5)}")
  721. if scheduler is not None:
  722. scheduler.step()
  724. if verbose_epochs:
  725. print("Training complete. Best parameters found at epoch", best_epoch)
  727. self.FSC.psi = best_psi
  728. for idx, param in enumerate(self.FSC.GPModel.param_names):
  729. self.FSC.GPModel.__setattr__(param, nn.Parameter(best_params[idx]))
  731. if NVal != 0:
  732. return losses_train, losses_val
  733. else:
  734. return losses_train, None
  736. def optimize(self, inference_params, verbose, verbose_epochs):
  737. """
  738. Run the full inference pipeline according to ``inference_params``.
  740. Depending on the configuration, this method executes MAPSO,
  741. gradient-based optimization, or both in sequence. It aggregates loss
  742. histories, computes ``best_loss``, and sets ``self.trained=True``.
  744. Parameters:
  745. --- inference_params: dict
  746. Validated inference configuration produced by
  747. ``FSC.set_inference_params``.
  748. --- verbose: bool
  749. Verbosity flag forwarded to optimization backends.
  750. --- verbose_epochs: bool
  751. If True, print per-epoch optimization updates.
  753. Returns:
  754. --- dict
  755. Loss history dictionary containing at least ``train`` and
  756. optionally ``MAPSO`` and ``val``.
  757. """
  759. assert self.trajectories_loaded, "No trajectories have been loaded. Load trajectories with the load_trajectories method."
  760. assert not self.trained, "The model has already been trained. If you want to train it again, reinitialize it or set the flag self.trained to False."
  762. loss_epochs = {}
  764. trainable_params = inference_params["trainable_parameters"]
  765. trainable_mask = inference_params["trainable_mask"]
  767. # check if any value of the trainable_mask dictionary is not None
  768. any_masked = False
  769. for key, val in trainable_mask.items():
  770. if val is not None:
  771. any_masked = True
  772. break
  774. if inference_params["use_MAPSO"]:
  775. n_particles = inference_params['n_particles_MAPSO']
  776. NEpochs = inference_params['NEpochs_MAPSO']
  778. c1_init = inference_params['c1_init_MAPSO']
  779. c2_init = inference_params['c2_init_MAPSO']
  780. w_init = inference_params['w_init_MAPSO']
  782. sigma_min = inference_params['sigma_min_MAPSO']
  783. sigma_max = inference_params['sigma_max_MAPSO']
  785. dynamic_topology = inference_params['dynamic_topology_MAPSO']
  786. num_neighbors_init = inference_params['num_neighbors_init_MAPSO']
  787. num_neighbors_final = inference_params['num_neighbors_final_MAPSO']
  788. num_neighbors_mid = inference_params['num_neighbors_mid_MAPSO']
  790. init_particles = inference_params['init_particles_MAPSO']
  791. init_velocities = inference_params['init_velocities_MAPSO']
  794. loss_MAPSO = self.optimize_w_MAPSO(trainable_params, trainable_mask,
  795. n_particles, NEpochs,
  796. init_particles, init_velocities,
  797. c1_init, c2_init, w_init,
  798. sigma_min, sigma_max,
  799. dynamic_topology,
  800. num_neighbors_init, num_neighbors_final, num_neighbors_mid,
  801. verbose, verbose_epochs)
  802. loss_epochs["MAPSO"] = loss_MAPSO
  804. if inference_params["use_gradient"]:
  806. NEpochs = inference_params['NEpochs_gradient']
  807. NBatch = inference_params['NBatch_gradient']
  808. lr = inference_params['lr_gradient']
  810. train_split = inference_params['train_split_gradient']
  811. scheduler = inference_params['scheduler_gradient']
  812. optimizer = inference_params['optimizer_gradient']
  814. if inference_params["use_ccopt"]:
  815. if self.FSC._init_memory_obs_dependent:
  816. raise ValueError("CCOpt cannot be used with memory observation dependent initialization.")
  818. maxiter = inference_params['maxiter_ccopt']
  819. rho0 = inference_params['rho0_ccopt']
  820. th = inference_params['th_ccopt']
  821. c_gauge = inference_params['c_gauge_ccopt']
  822. else:
  823. maxiter = None
  824. rho0 = None
  825. th = None
  826. c_gauge = None
  828. losses_gradient = self.optimize_w_gradient(inference_params["use_ccopt"],
  829. trainable_params, trainable_mask, any_masked,
  830. NEpochs, NBatch, lr,
  831. train_split, optimizer, scheduler,
  832. maxiter, rho0, th, c_gauge,
  833. verbose, verbose_epochs)
  835. if inference_params["use_MAPSO"]:
  836. loss_epochs["train"] = np.concatenate([loss_epochs["MAPSO"], losses_gradient[0][1:]])
  837. if train_split < 1.0:
  838. loss_epochs["val"] = np.concatenate([loss_epochs["MAPSO"], losses_gradient[1][1:]])
  839. self.best_loss = np.min(loss_epochs["val"])
  840. else:
  841. self.best_loss = np.min(loss_epochs["train"])
  842. else:
  843. loss_epochs["train"] = losses_gradient[0]
  844. if train_split < 1.0:
  845. loss_epochs["val"] = losses_gradient[1]
  846. self.best_loss = np.min(loss_epochs["val"])
  847. else:
  848. self.best_loss = np.min(loss_epochs["train"])
  849. else:
  850. loss_epochs["train"] = loss_MAPSO
  851. self.best_loss = np.min(loss_epochs["train"])
  853. self.trained = True
  855. return loss_epochs
  857. def get_inferred_policy_params(self):
  858. """
  859. Return the current inferred policy parameters.
  861. Returns:
  862. --- list of torch.Tensor
  863. Policy parameters in the order of
  864. ``self.FSC.GPModel.param_names``.
  865. """
  866. return [self.FSC.GPModel.__getattribute__(param) for param in self.FSC.GPModel.param_names]
  869. @staticmethod
  870. @nb.njit
  871. def optimize_rho(Y, M, A, TMat, pya, rhok, maxiter, th):
  872. """
  873. Numba-accelerated convex-concave fixed-point update for rho.
  875. Given the current transition matrix and empirical start distribution
  876. over observation-action pairs, iteratively updates ``rho`` until
  877. convergence or the maximum number of iterations is reached.
  879. Parameters:
  880. --- Y: int
  881. Number of observations.
  882. --- M: int
  883. Number of memory states.
  884. --- A: int
  885. Number of actions.
  886. --- TMat: np.ndarray
  887. Transition tensor with shape ``(Y, A, M, M)``.
  888. --- pya: np.ndarray
  889. Empirical start distribution over ``(y, a)`` with shape ``(Y, A)``.
  890. --- rhok: np.ndarray
  891. Current rho iterate with shape ``(M,)``.
  892. --- maxiter: int
  893. Maximum number of fixed-point iterations.
  894. --- th: float
  895. Convergence threshold on ``||rho_{k+1} - rho_k||``.
  897. Returns:
  898. --- tuple
  899. ``(rho, err)`` where ``rho`` is the final iterate and ``err`` is
  900. the final norm difference between consecutive iterates.
  901. """
  902. TMat = np.transpose(TMat, (0, 2, 3, 1))
  903. wVec = np.zeros((Y, A, M))
  904. for y in range(Y):
  905. for a in range(A):
  906. for m in range(M):
  907. wVec[y, a, m] = np.sum(TMat[y, m, :, a])
  909. for _ in range(maxiter):
  910. wsumexp_test_k = np.zeros((Y, A))
  912. for y in range(Y):
  913. for a in range(A):
  914. wsumexp_test_k[y, a] = np.sum(wVec[y, a] * rhok)
  916. grad = wVec * rhok / wsumexp_test_k[..., None]
  918. rhok_new = np.zeros(M)
  920. for y in range(Y):
  921. for a in range(A):
  922. rhok_new += pya[y, a] * grad[y, a]
  924. if np.linalg.norm(rhok_new - rhok) < th:
  925. break
  927. rhok = rhok_new
  929. return rhok, np.linalg.norm(rhok_new - rhok)

inference.py at commit 219eb08, no license · at the source

Overview

Authors: Giorgio Nicoletti1, Antonio Celani1,2
ORCID iDs: Giorgio Nicoletti, Antonio Celani
  1. Quantitative Life Sciences section, The Abdus Salam International Center for Theoretical Physics (ICTP), Trieste, Italy
  2. Department of Oncology, Università degli Studi di Torino, Italy
Institutions: The Abdus Salam International Centre for Theoretical Physics (ICTP) (Italy); University of Turin (Italy)
Journal: PLoS computational biology, volume 22, issue 8, article e1014585
Dates: received 17 February 2026; accepted 16 July 2026; published online 7 August 2026
Type: Research article · Language: English
License: CC BY
Identifiers: DOI 10.1371/journal.pcbi.1014585 · PMID 42566485 · PMCID PMC13475988 · OpenAlex W7201868829
Open access: gold, a free copy (OpenAlex)
Status: code verified
Categories: computational modeling (no new data) (modality), mouse (organism), rat (organism), cognitive (subfield)
Methods: Machine learning, fMRI & imaging
MeSH: Behavior, Animal*, Decision Making*, Models, Neurological*, Animals, Computational Biology, Computer Simulation, Mice, Rats (* major topic)
Topic: Zebrafish Biomedical Research Applications (Cell Biology, Biochemistry, Genetics and Molecular Biology), according to OpenAlex
Citations: not cited yet (Europe PMC); 62 references in the paper

Abstract

Understanding how living organisms process sensory information from their surroundings and translate it into decisions is a fundamental problem across biological scales – from biochemical signalling in single-cells to neural computations in animal brains. In this work, we address this challenge by introducing a method to reconstruct general decision processes directly from behavioral observations alone. Our approach is applicable to any biological agent and does not require prior knowledge of its internal mechanisms or its environment. Our agent model is defined by a recurrent dynamics over a discrete set of internal states which encode and process sensory information, and dictate which actions to execute. We validate our method on synthetic agents and demonstrate that we can exactly recover the agent’s behavior for non-trivial tasks. Then, we infer agent models from experimental data of rats performing evidence accumulation and of mice making decisions under uncertainty and in changing environments. In both cases, very few internal states suffice to reproduce the observed behavior with high accuracy. Crucially, the immediate interpretability of the inferred dynamics allows to understand the computational process underlying decision-making.

Reproduced under the paper's license (CC BY), from the paper cited above.

Repository

Its files are read in the Code ↔ Paper reader above, with 6 matches between paragraphs and lines of code.

giorgionicoletti/FSC-inference-MAPSO

License: none: the authors keep all their rights
State: the link answers, verified on 27 September 2026
Evidence: files inventoried
Commit: 219eb08e8f79777470dbcd0896eaa035cb693c2f, 15 May 2026
Languages: Python (13), Jupyter (5)
Size: 260 files, 18 scripts
Software Heritage: not archived
Found in: “Data Availability”
Holds: README, 5 notebooks
Not found: license file, CITATION.cff, environment file, tests, continuous integration, documentation
Tools: NumPy (16 files), Numba (9 files), Matplotlib (8 files), PyTorch (5 files), pandas (3 files), NetworkX (1 file), SciPy (1 file)
Availability: 1 check, the latest on 27 September 2026: the link answers
  • 27 September 2026: the link answers
19 files
At the source: github.com/giorgionicoletti/FSC-inference-MAPSO

The paper's code and data availability statement is in the Data section.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 1 repository of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 18 scripts, each with its path and the digest of its content;
  • 6 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

No dataset and no data link were found in the paper.

Data Availability

The code to infer finite state controllers from behavioral trajectories is available at https://github.com/giorgionicoletti/FSC-inference-MAPSO.

Reproduced under the paper's license (CC BY), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 27 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 2 authors, 8 MeSH terms, 51 references.

Cite

This paper

Nicoletti, G., & Celani, A. (2026). Decoding behavior with minimal and interpretable agent models. PLoS computational biology, 22(8), e1014585. https://doi.org/10.1371/journal.pcbi.1014585

BibTeX

@article{nicoletti2026decoding,
author = {Nicoletti, Giorgio and Celani, Antonio},
title = {{Decoding behavior with minimal and interpretable agent models}},
journal = {PLoS computational biology},
year = {2026},
month = aug,
volume = {22},
number = {8},
pages = {e1014585},
publisher = {PLOS},
issn = {1553-734X},
doi = {10.1371/journal.pcbi.1014585},
url = {https://doi.org/10.1371/journal.pcbi.1014585},
pmid = {42566485},
pmcid = {PMC13475988}
}

RIS

TY - JOUR
AU - Nicoletti, Giorgio
AU - Celani, Antonio
TI - Decoding behavior with minimal and interpretable agent models
T2 - PLoS computational biology
J2 - PLoS Comput Biol
PY - 2026
DA - 2026/08/07
VL - 22
IS - 8
SP - e1014585
SN - 1553-734X
PB - PLOS
DO - 10.1371/journal.pcbi.1014585
UR - https://doi.org/10.1371/journal.pcbi.1014585
LA - en
ER -

CSL-JSON

{
"id": "10.1371/journal.pcbi.1014585",
"type": "article-journal",
"title": "Decoding behavior with minimal and interpretable agent models",
"container-title": "PLoS computational biology",
"author": [
{
"family": "Nicoletti",
"given": "Giorgio"
},
{
"family": "Celani",
"given": "Antonio"
}
],
"container-title-short": "PLoS Comput Biol",
"volume": "22",
"issue": "8",
"page": "e1014585",
"DOI": "10.1371/journal.pcbi.1014585",
"PMID": "42566485",
"PMCID": "PMC13475988",
"ISSN": "1553-734X",
"publisher": "PLOS",
"URL": "https://doi.org/10.1371/journal.pcbi.1014585",
"language": "en",
"issued": {
"date-parts": [
[
2026,
8,
7
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.1038/s41467-026-75924-7 [code]
Data-driven reduced modeling of neural dynamics.
Journal: Nature communications
In common: SciPy, Matplotlib, NumPy, computational modeling (no new data), 6 references
[2] doi:10.3390/biomimetics11080569 [code]
Pretraining of Embodied Recurrent Networks Bridges the Gap Between Artificial and Cortical Neural Activities.
Journal: Biomimetics (Basel, Switzerland)
In common: PyTorch, SciPy, Matplotlib, 1 other tool, 5 references
[3] doi:10.1038/s41467-026-72057-9 [code]
Sex-specific behavioral feedback modulates sensorimotor processing and drives flexible social behavior.
Journal: Nature communications
In common: Numba, PyTorch, pandas, 3 other tools, 3 references
[4] doi:10.1371/journal.pbio.3003831 [code]
Disinhibitory signaling enables flexible coding of top-down information in cortical networks.
Journal: PLoS biology
In common: NetworkX, PyTorch, pandas, 3 other tools, mouse, 3 references
[5] doi:10.7554/elife.109313 [code]
Linear and categorical coding units in the mouse gustatory cortex drive population dynamics and behavior in taste decision-making.
Journal: eLife
In common: Numba, PyTorch, pandas, 3 other tools, computational modeling (no new data), cognitive, mouse, 1 reference
[6] doi:10.1016/j.celrep.2026.117262 [code]
Cerebellar acceleration of learning in an evidence-accumulation task.
Journal: Cell reports
In common: pandas, SciPy, Matplotlib, 1 other tool, cognitive, mouse, 3 references
[7] doi:10.1016/j.isci.2026.116825 [code]
Social hierarchy shapes behavioral and transcriptional responses to chronic stress and ketamine in male mice.
Journal: iScience
In common: Numba, NetworkX, PyTorch, 4 other tools, mouse, 1 reference
[8] doi:10.1038/s41593-026-02232-0 [code]
Entorhinal cortex represents task-relevant remote locations independently of CA1.
Journal: Nature neuroscience
In common: Numba, NetworkX, PyTorch, 4 other tools, mouse, 1 reference
[9] doi:10.3389/fnsys.2026.1822122 [code]
Convergence-divergence circuits for multimodal integration of innate and learned opponent valences.
Journal: Frontiers in systems neuroscience
In common: Numba, NetworkX, PyTorch, 4 other tools, 1 reference
[10] doi:10.1371/journal.pcbi.1014162 [code]
Exploring neural manifolds across a wide range of intrinsic dimensions.
Journal: PLoS computational biology
In common: pandas, SciPy, Matplotlib, 1 other tool, 4 references

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.