env
===

.. py:module:: env

.. autoapi-nested-parse::

   Custom mRNA codon design environment using torchgfn to generate mRNA sequences encoding a given protein. It supports a multi-objective optimization over biological properties of mRNA sequences. Implemented using the DiscreteEnv class.
   Each timestep corresponds to choosing a synonymous codon for the next amino acid in the sequence.
   **Action Space**:
   Number of CODONS + 1 possible actions (all codons + 1 exit action)
   **State Representation**:
   A vector of length = protein length, initialized to -1. Codons are filled in step-by-step.
   Masking (Action Constraints)**:
   At each position t, only codons that correspond to the t-th amino acid are allowed to ensure biological correctness.
   **Reward function**:
   A combination of multiple biological properties to evaluate the mRNA sequence. Weights of these objectives can be updated dynamically to reflect different reward configurations. Rewards and constraints are modular and can be extended to incorporate new objectives.
   The environment is customizable for different organisms by using species-specific codon tables preferences, and could serve as a benchmark environment in computational biology. This enables exploration of codon space, which is a large search space given a protein sequence, to optimize for mRNA design.  Applicable to mRNA vaccines, protein therapeutics, and gene expression optimization.



Classes
-------

.. autoapisummary::

   env.CodonDesignEnv


Module Contents
---------------

.. py:class:: CodonDesignEnv(protein_seq, device, sf=None)

   Bases: :py:obj:`gfn.env.DiscreteEnv`


   Environment for designing mRNA codon sequences for a given protein.
   Action space is the global codon set (size N_CODONS) plus an exit action.
   Dynamic masks restrict actions at each step:
   - At step t < seq_length: only synonymous codons for protein_seq[t] are allowed.
   - At step t == seq_length: only the exit action is allowed.



   .. py:attribute:: States
      :type:  type[gfn.states.DiscreteStates]


   .. py:attribute:: _device


   .. py:method:: backward_step(states, actions)

      Backward transition function of the environment.

      This method takes a batch of states and actions and returns a batch of previous
      states. It does not need to check whether the actions are valid or the states are
      sink states, because the `_backward_step` method wraps it and checks for validity.

      :param states: A batch of states.
      :param actions: A batch of actions.

      :returns: A batch of previous states.



   .. py:attribute:: codon_gc_counts


   .. py:attribute:: exit_action_index


   .. py:attribute:: idx_to_codon


   .. py:method:: is_terminal(states)


   .. py:method:: make_sink_states_tensor(shape, device=None)
      :staticmethod:



   .. py:attribute:: n_actions


   .. py:attribute:: protein_seq


   .. py:method:: reward(final_states)

      Returns the environment's rewards for a batch of states.

      This or `log_reward` must be implemented by the environment.

      :param states: A batch of states with a batch_shape.

      :returns: Tensor of shape (*batch_shape) containing the rewards.



   .. py:attribute:: seq_length


   .. py:method:: set_weights(w)

      Store the current preference weights (w) for conditional reward.



   .. py:method:: step(states, actions)

      Forward transition function of the environment.

      This method takes a batch of states and actions and returns a batch of next
      states. It does not need to check whether the actions are valid or the states are
      sink states, because the `_step` method wraps it and checks for validity.

      :param states: A batch of states.
      :param actions: A batch of actions.

      :returns: A batch of next states.



   .. py:attribute:: syn_indices


   .. py:method:: update_masks(states)


   .. py:attribute:: weights