fix reshaping in predict next state

lasgroup · Aug 30, 2023 · a0f24de · a0f24de
1 parent a2d3ea6
commit a0f24de
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/experiments/model_based_rl_sim_transfer_comparison/exp.py b/experiments/model_based_rl_sim_transfer_comparison/exp.py
@@ -199,6 +199,6 @@ def main(args):
     parser.add_argument('--best_bnn_model', type=int, default=1)
     parser.add_argument('--best_policy', type=int, default=0)
     parser.add_argument('--margin_factor', type=float, default=20.0)
-    parser.add_argument('--predict_difference', type=int, default=1)
+    parser.add_argument('--predict_difference', type=int, default=0)
     args = parser.parse_args()
     main(args)
diff --git a/experiments/model_based_rl_sim_transfer_comparison/launcher.py b/experiments/model_based_rl_sim_transfer_comparison/launcher.py
@@ -1,7 +1,7 @@
 import exp
 from experiments.util import generate_run_commands, generate_base_command, dict_permutations
 
-PROJECT_NAME = 'PredictNextStateInsteadofDifference'
+PROJECT_NAME = 'PredictNextStateInsteadofDifferenceN2'
 
 applicable_configs = {
     'horizon_len': [50, 2 ** 6, 100],
@@ -10,7 +10,7 @@
     'num_episodes': [40],
     'sac_num_env_steps': [1_000_000, 2_000_000],
     'bnn_train_steps': [20_000, 40_000],
-    'learnable_likelihood_std': ['yes', 'no'],
+    'learnable_likelihood_std': ['yes'],
     'reset_bnn': ['no'],
     'use_sim_prior': [1],
     'include_aleatoric_noise': [1],

diff --git a/sim_transfer/rl/model_based_rl/learned_system.py b/sim_transfer/rl/model_based_rl/learned_system.py
@@ -46,6 +46,7 @@ def next_state(self,
             x_next_dist = self.model.predict_dist(z, include_noise=self.include_noise)
             next_key, key_sample_x_next = jr.split(dynamics_params.key)
             x_next = x_next_dist.sample(seed=key_sample_x_next)
+            x_next = x_next.reshape((self.x_dim,))
         new_dynamics_params = dynamics_params.replace(key=next_key)
         return Normal(loc=x_next, scale=jnp.zeros_like(x_next)), new_dynamics_params