Factor reshape to function

[krzysztofrusek]

ltc/ltc/symbolic/history2csv.py

Lines 73 to 83 in e8c78ca

	observations = history.observations[-1, ...]
	n_steps, _, window_size, _ = observations.shape
	key = jax.random.key(42)
	qvals = compute_qvals(
	params, state, observations, key
	) # [n_agents, n_steps, n_actions]
	# Build flat feature matrix: [n_agents * n_steps, window_size * n_features]
	# transpose [n_steps, n_agents, w, f] -> [n_agents, n_steps, w, f] then flatten last two dims
	XX = np.asarray(observations).transpose(1, 0, 2, 3).reshape(n_agents * n_steps, -1)