Reinforcement learning (2405.10369v1)

Abstract: Observing celestial objects and advancing our scientific knowledge about them involves tedious planning, scheduling, data collection and data post-processing. Many of these operational aspects of astronomy are guided and executed by expert astronomers. Reinforcement learning is a mechanism where we (as humans and astronomers) can teach agents of artificial intelligence to perform some of these tedious tasks. In this paper, we will present a state of the art overview of reinforcement learning and how it can benefit astronomy.

• The paper presents an overview of reinforcement learning fundamentals and its deep learning extensions tailored for astronomy.
• It details model-free, model-based, and hint-assisted RL techniques for optimizing telescope control and data processing.
• The work highlights practical challenges such as state design and reward shaping, offering source code for real-world implementations.

This paper provides an overview of reinforcement learning (RL) with a focus on its potential applications in astronomy. It aims to equip astronomers with the foundational knowledge needed to apply modern deep RL techniques to tasks like telescope automation, observation scheduling, and data processing hyper-parameter tuning.

1. Reinforcement Learning Fundamentals

The core RL framework involves an agent interacting with an environment.

• The agent observes the environment's state ($s \in \mathcal{S}$).
• Based on the state, the agent takes an action ($a \in \mathcal{A}$).
• The environment transitions to a new state ($s'$) and provides a scalar reward ($r \in \mathcal{R}$) to the agent.
• The goal is to learn a policy ($\pi$) that maximizes the cumulative discounted reward over time.

This interaction is often modeled as a Markov Decision Process (MDP), defined by the tuple $$(\mathcal{S}, \mathcal{A}, \mathcal{R}, \mathcal{P})$$, where $\mathcal{P}$ represents the state transition probabilities $p(s'|s,a)$. Key concepts include:

• Q-function $Q(s,a)$: The expected cumulative reward starting from state $s$, taking action $a$, and following the policy thereafter.
• Value function $V(s)$: The expected cumulative reward starting from state $s$ and following the policy.
• Policy $\pi(s) \rightarrow a$ (deterministic) or $\pi(a|s)$ (stochastic): Maps states to actions or action probabilities.

The optimal Q-function and policy are related by the BeLLMan equation:

$Q(s,a) = r(s,a) + \gamma \max_{a'} Q(s', a')$

where $\gamma$ is the discount factor. For simple, discrete problems, this can be solved iteratively using a Q-table (demonstrated with a maze example). For complex problems with high-dimensional or continuous states/actions, functions are approximated using Deep Neural Networks (DNNs).

2. Deep RL Algorithms (Model-Free)

These algorithms learn directly from interactions with the environment without explicitly modeling its dynamics.

• Challenges: Data inefficiency, balancing exploration (trying new actions) vs. exploitation (using known good actions), and training stability.
• Experience Replay: Storing past transitions $(s, a, r, s')$ in a replay buffer $\mathcal{D}$ and sampling mini-batches from it improves data efficiency and stability. This necessitates off-policy algorithms.
• General Training Loop: Involves iterating through episodes, where the agent selects actions, interacts with the environment, stores experiences, and periodically samples from the buffer to update its networks (Algorithm 2).

Algorithms for Discrete Actions:

• Q-learning: Iteratively updates the Q-value estimate using the BeLLMan equation (Eq. 3).
• Double Q-learning: Uses two Q-networks ($Q_1, Q_2$) to decouple action selection and value estimation, reducing overestimation bias (Eqs. 4-5).
• Deep Q-Network (DQN): Represents Q-functions with DNNs ($Q_\theta$). Uses a target network ($Q_{\theta'}$) for stability, minimizing the mean squared error loss (Eq. 6) via gradient descent (Eq. 7). The target network parameters $\theta'$ are periodically updated with $\theta$.

Algorithms for Continuous Actions (Actor-Critic): These methods maintain separate networks for the policy (actor) and the value/Q-function (critic).

• Deep Deterministic Policy Gradient (DDPG): Learns a deterministic policy $\pi_\phi(s)$. Uses target networks for both actor ($\pi_{\phi'}$) and critic ($Q_{\theta'}$). The critic is updated by minimizing TD error (Eq. 9), and the actor is updated by maximizing the expected Q-value (minimizing $-Q_\theta(s, \pi_\phi(s))$, Eq. 10). Target networks are updated using Polyak averaging (Eq. 12). Action noise (e.g., Ornstein-Uhlenbeck) is added for exploration (Eq. 8).
• Twin Delayed DDPG (TD3): Improves DDPG by:
  • Using two critic networks ($Q_{\theta_1}, Q_{\theta_2}$) and taking the minimum of their target values to mitigate Q-value overestimation (Eq. 16).
  • Delaying policy and target network updates relative to critic updates.
  • Adding clipped noise to the target policy action and clipping the resulting action for target policy smoothing (Eqs. 14-15).
• Soft Actor-Critic (SAC): Learns a stochastic policy $\pi_\phi(a|s)$ and aims to maximize both cumulative reward and policy entropy (encouraging exploration).
  • Uses twin critics and target critics similar to TD3.
  • The critic loss includes an entropy term $-\alpha \log \pi_\phi(a'|s')$ (Eq. 20).
  • The actor loss also includes the entropy term (Eq. 22).
  • Uses the reparameterization trick for sampling actions differentiably (Eq. 24).

3. Model-Based RL

These methods learn a model of the environment's dynamics $p(s'|s,a)$ and use it for planning or generating synthetic data.

• Motivation: Improved data efficiency and safety, as fewer real-world interactions are needed.
• Probabilistic Ensembles: Uses an ensemble of $B$ probabilistic DNNs to model the dynamics $$p_{\theta_i}(s'|s,a) \sim \mathcal{N}(\mu_{\theta_i}, \Sigma_{\theta_i})$$. This captures both aleatoric (inherent randomness) and epistemic (model uncertainty) uncertainty. Each model is trained by minimizing negative log-likelihood on different bootstrapped samples from the replay buffer (Eq. 26).
• Probabilistic Ensemble with Trajectory Sampling (PETS): Uses the learned ensemble model for planning. At each step $t$, it employs the Cross-Entropy Method (CEM, Algorithm 4) to find the optimal action $a_t$. CEM samples action sequences, simulates trajectories using randomly chosen models from the ensemble, evaluates the expected rewards, and iteratively refines the action distribution towards high-reward sequences (Algorithm 3).

4. Hint Assisted RL

This approach incorporates existing knowledge (e.g., from heuristics, simpler models, or human experts) into the RL training process.

• A hint $h$ (representing a suggested action) is provided to the actor.
• A constraint $c(a,h)$ measures the distance between the actor's action $a$ and the hint $h$.
• The policy optimization objective is augmented using the Alternating Direction Method of Multipliers (ADMM) to encourage the policy action $a_\phi$ to stay close to the hint $h$, controlled by a threshold $\delta$ (Eqs. 27-30). This allows leveraging prior knowledge without strictly enforcing it.

5. Applications in Astronomy & Practical Considerations

• Potential Applications: Telescope control (adaptive optics, scheduling), resource allocation (compute, observing time), hyper-parameter tuning in data processing pipelines, and new science discovery from archival data.
• Implementation Issues:
  • State/Action Design: Requires domain knowledge and experimentation. May need to include history if the Markov property doesn't hold.
  • Normalization: Crucial for DNN stability when combining heterogeneous data.
  • Reward Shaping: Designing effective reward functions is key; can involve scaling, penalties, and clipping.
  • Hybrid Actions: Combining discrete and continuous actions requires specific techniques (e.g., predicting probabilities for discrete parts).
  • Variable Dimensions: Techniques like padding, auto-encoders, or attention mechanisms can handle inputs/outputs of varying sizes (e.g., sky models).

6. Examples

• Bipedal Walker: A standard benchmark. SAC outperforms TD3 on the simple version. Hint-assisted SAC and TD3 show improved performance on the difficult "hardcore" version, demonstrating the benefit of incorporating knowledge (hints from an agent trained on the easy version).
• Calibration Example: Formulates a model selection/fitting problem (Eq. 31) as an RL task.
  • Objective: Select basis functions and computational budget to optimize fit quality (AIC) and cost.
  • State: Based on influence functions summarizing model performance.
  • Action: Vector predicting inclusion probability for each basis function and a scaled budget.
  • Reward: Based on negative AIC and computational penalty (Eqs. 32-34).
  • Results show a SAC agent learning the task, benefiting from hints derived from exhaustive search.

Conclusion

Deep RL offers powerful tools for automating complex tasks in astronomy. Model-free, model-based, and hint-assisted approaches provide flexibility. While practical implementation requires careful consideration of state/action/reward design and network tuning, the potential benefits for observatory operations, data analysis, and scientific discovery are significant. The paper provides source code for the discussed algorithms.